diff --git a/README.md b/README.md index eba4306..07f302b 100644 --- a/README.md +++ b/README.md @@ -106,7 +106,7 @@ python -m rrparser --dataset rhea --diameters '<=8' -o rules.csv python -m rrparser --dataset rhea --diameters '>=8' -o rules.csv ``` -`--radii` supports the same expressions over integer radii 0–10: +`-r` / `--radii` supports the same expressions over integer radii 0–10: `0,1,2`, `0-4`, `<4`, `<=4`, `>4`, `>=4`, and combinations such as `0-2,>=8`. For example, `--radii '<=4'` is equivalent to `--diameters '<=8'`. The two options are mutually exclusive. Both are expanded into normalized @@ -114,8 +114,11 @@ integer radii by the argument parser, before calling `parse_rules()`. Without either option, all radii 0–10 (diameters 0,2,…,20) are selected. For templates, a row is retained when at least one source radius matches -`diameter / 2`, once per template. The output `RADII` column retains all source -radii, separated by semicolons. Explicit odd diameters or values above 20 are +`diameter / 2`, once per template. The output `RADII` column retains only the +selected radii, separated by semicolons; `RADIUS_MIN` and `RADIUS_MAX` are +updated to match. For example, `-r '>4'` exports only radii 5–10 even if a +matching source template also covers radius 4. Cached sources are unchanged. +Explicit odd diameters or values above 20 are rejected before parsing, including for rr02. rr02 supplies diameters 2,4,…,16. The extended default therefore selects all @@ -161,7 +164,7 @@ archive cache. | `--list-releases` | Print available release/dataset pairs and exit. | | `--cache-dir`, `--rules-dir` | Source cache directory; default package-local `rrparser/cache`. | | `-d`, `--diameters` | List, range, or comparison; default all even diameters 0–20. | -| `--radii` | List, range, or comparison over radii 0–10, as an alternative to `--diameters`. | +| `-r`, `--radii` | List, range, or comparison over radii 0–10, as an alternative to `--diameters`. | | `-rt`, `--rule-type` | `all` (default), `retro`, or `forward`; templates require `all`. | | `-ec`, `--ec` | File of comma-separated EC prefixes to keep. | | `-ecx`, `--ecx` | File of comma-separated EC prefixes to exclude; mutually exclusive with `--ec`. | diff --git a/rrparser/Args.py b/rrparser/Args.py index 431912a..3902ea3 100644 --- a/rrparser/Args.py +++ b/rrparser/Args.py @@ -142,6 +142,7 @@ def _add_arguments(parser): "For templates, diameter = 2 * radius; rr02 provides diameters 2 to 16", ) sizes.add_argument( + "-r", "--radii", type=parse_radii, help='Radii: list (0,1,2), inclusive range (0-4), "<4", "<=4", ">4" or ">=4". ' diff --git a/rrparser/Parser.py b/rrparser/Parser.py index 1b228a3..376d24c 100644 --- a/rrparser/Parser.py +++ b/rrparser/Parser.py @@ -320,16 +320,18 @@ def filter_( raise ValueError( "RetroRules templates have no retro/forward annotation; use --rule-type all" ) - radius_idx = df[ + selection = set(radii) + selected_radii = ( df["RADII"] .fillna("") .astype(str) - .apply( - lambda values: any( - int(radius) in radii for radius in values.split(";") if radius + .map( + lambda values: sorted( + {int(radius) for radius in values.split(";") if radius} & selection ) ) - ].index + ) + radius_idx = df[selected_radii.map(bool)].index query_vars["radius_idx"] = radius_idx query += " & index in @radius_idx" else: @@ -340,6 +342,13 @@ def filter_( logger.debug(f"Query: {query}") try: _df = df.query(query, local_dict=query_vars) + if "RADII" in _df.columns: + _df = _df.copy() + retained_radii = selected_radii.loc[_df.index] + _df["RADII"] = retained_radii.map(lambda values: ";".join(map(str, values))) + for column, bound in (("RADIUS_MIN", min), ("RADIUS_MAX", max)): + if column in _df.columns: + _df[column] = retained_radii.map(bound) logger.info("") logger.info(f"=> {len(_df)} rules found") logger.info("") diff --git a/tests/test_aggregation.py b/tests/test_aggregation.py index 74eb23c..ebbe403 100644 --- a/tests/test_aggregation.py +++ b/tests/test_aggregation.py @@ -80,6 +80,9 @@ def test_cached_datasets_are_merged_before_filtering(tmp_path, selection): assert shared["RADIUS_MIN"] == 0 assert shared["RADIUS_MAX"] == 10 assert pd.read_csv(StringIO(filtered))["Rule ID"].tolist() == ["shared"] + filtered_row = pd.read_csv(StringIO(filtered), dtype={"RADII": str}).iloc[0] + assert filtered_row["RADII"] == "10" + assert filtered_row["RADIUS_MIN"] == filtered_row["RADIUS_MAX"] == 10 assert [source.read_bytes() for source in sources] == originals diff --git a/tests/test_radii.py b/tests/test_radii.py index a41b8e2..4b25467 100644 --- a/tests/test_radii.py +++ b/tests/test_radii.py @@ -50,6 +50,7 @@ def test_radius_expressions_select_templates(tmp_path, expression, expected): [ ("--radii", "0-4", [0, 1, 2, 3, 4]), ("--radii", "<=4", [0, 1, 2, 3, 4]), + ("-r", ">4", [5, 6, 7, 8, 9, 10]), ("--diameters", "0-8", [0, 1, 2, 3, 4]), ("-d", "<=8", [0, 1, 2, 3, 4]), ("--radii", ">=8", [8, 9, 10]), @@ -108,3 +109,43 @@ def test_defaults_and_empty_selection(): args = build_args_parser().parse_args([]) assert list(args.radii) == list(range(11)) assert parse_diameters("<0") == parse_radii("<0") == [] + + +@pytest.mark.parametrize( + "option,expression,expected", + [ + ("-r", ">4", "5;6;7;8;9;10"), + ("--radii", ">=4", "4;5;6;7;8;9;10"), + ("-r", "4-6", "4;5;6"), + ("-r", "0,10", "10"), + ("-d", ">8", "5;6;7;8;9;10"), + ("-r", ">10", ""), + ], +) +def test_exported_radius_metadata_matches_selection( + tmp_path, option, expression, expected +): + source = tmp_path / "templates.csv" + pd.DataFrame( + { + "TEMPLATE_ID": ["RR:shared"], + "TEMPLATE": ["[C]>>[O]"], + "SCORE": [1.0], + "ECS": ["1.1.1.1"], + "RADII": ["4;5;6;7;8;9;10"], + "RADIUS_MIN": [4], + "RADIUS_MAX": [10], + } + ).to_csv(source, index=False) + original = source.read_bytes() + args = build_args_parser().parse_args([option, expression]) + result = parse_rules(None, rules_file=str(source), radii=args.radii) + output = pd.read_csv(StringIO(result), dtype={"RADII": str}) + if expected: + assert output["RADII"].tolist() == [expected] + radii = list(map(int, expected.split(";"))) + assert output["RADIUS_MIN"].tolist() == [min(radii)] + assert output["RADIUS_MAX"].tolist() == [max(radii)] + else: + assert output.empty + assert source.read_bytes() == original