From cdf7f3752ca533aacbaa5dc2d152907fb0b8c8e1 Mon Sep 17 00:00:00 2001 From: Victor Lin Date: Fri, 11 Sep 2026 15:00:47 -0700 Subject: [PATCH 1/4] Allow $schema key in command config files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit $schema is used in the schema to refer to the meta-schema,¹ but it is also common practice to use the same key in schema-adherent files for strictness, similar to typing a variable in code. One benefit is IDE hints² for users who open these config files in a editor that understands this key, but more importantly we can use it in our own processing of config files. ¹ https://json-schema.org/understanding-json-schema/reference/schema#schema ² https://code.visualstudio.com/docs/languages/json#_json-schemas-and-settings --- augur/argparse_.py | 1 + augur/data/schema-align-config.json | 3 +++ augur/data/schema-ancestral-config.json | 3 +++ augur/data/schema-frequencies-config.json | 3 +++ augur/data/schema-mask-config.json | 3 +++ augur/data/schema-refine-config.json | 3 +++ augur/data/schema-subsample-config-unaligned.json | 3 +++ augur/data/schema-subsample-config.json | 3 +++ augur/data/schema-traits-config.json | 3 +++ augur/data/schema-translate-config.json | 3 +++ augur/data/schema-tree-config.json | 3 +++ devel/regenerate-config-schemas | 10 ++++++++-- devel/regenerate-subsample-schema | 3 +++ docs/_ext/config_schema_helpers.py | 4 ++++ 14 files changed, 46 insertions(+), 2 deletions(-) diff --git a/augur/argparse_.py b/augur/argparse_.py index 4653ca487..1c17ca610 100644 --- a/augur/argparse_.py +++ b/augur/argparse_.py @@ -143,6 +143,7 @@ def parse_known_args(self, args, namespace=None, **kwargs): try: with open(config_file) as f: config_keys = set(self._config_file_parser.parse(f).keys()) + config_keys.discard("$schema") except configargparse.ConfigFileParserException as err: self.error_only(str(err)) diff --git a/augur/data/schema-align-config.json b/augur/data/schema-align-config.json index 32675640b..ced27936a 100644 --- a/augur/data/schema-align-config.json +++ b/augur/data/schema-align-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/align-config/v1" + }, "sequences": { "description": "sequences to align", "oneOf": [ diff --git a/augur/data/schema-ancestral-config.json b/augur/data/schema-ancestral-config.json index 63876b383..f4ba9e32c 100644 --- a/augur/data/schema-ancestral-config.json +++ b/augur/data/schema-ancestral-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/ancestral-config/v1" + }, "tree": { "description": "prebuilt Newick", "type": "string", diff --git a/augur/data/schema-frequencies-config.json b/augur/data/schema-frequencies-config.json index 78eb2513a..7493ed47c 100644 --- a/augur/data/schema-frequencies-config.json +++ b/augur/data/schema-frequencies-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/frequencies-config/v1" + }, "method": { "description": "method by which frequencies should be estimated", "type": "string", diff --git a/augur/data/schema-mask-config.json b/augur/data/schema-mask-config.json index fd9cd63d0..9e53c4119 100644 --- a/augur/data/schema-mask-config.json +++ b/augur/data/schema-mask-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/mask-config/v1" + }, "sequences": { "description": "sequences in VCF or FASTA format", "type": "string", diff --git a/augur/data/schema-refine-config.json b/augur/data/schema-refine-config.json index ef550a923..be05e2650 100644 --- a/augur/data/schema-refine-config.json +++ b/augur/data/schema-refine-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/refine-config/v1" + }, "alignment": { "description": "alignment in fasta or VCF format", "type": "string", diff --git a/augur/data/schema-subsample-config-unaligned.json b/augur/data/schema-subsample-config-unaligned.json index e7ae57392..808b5fbde 100644 --- a/augur/data/schema-subsample-config-unaligned.json +++ b/augur/data/schema-subsample-config-unaligned.json @@ -306,6 +306,9 @@ } }, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/subsample-config-unaligned/v1" + }, "defaults": { "$ref": "#/$defs/defaultProperties" }, diff --git a/augur/data/schema-subsample-config.json b/augur/data/schema-subsample-config.json index b4e82979c..b08148a52 100644 --- a/augur/data/schema-subsample-config.json +++ b/augur/data/schema-subsample-config.json @@ -346,6 +346,9 @@ } }, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/subsample-config/v1" + }, "defaults": { "$ref": "#/$defs/defaultProperties" }, diff --git a/augur/data/schema-traits-config.json b/augur/data/schema-traits-config.json index 7e6c006d3..384f27de1 100644 --- a/augur/data/schema-traits-config.json +++ b/augur/data/schema-traits-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/traits-config/v1" + }, "tree": { "description": "tree to perform trait reconstruction on", "type": "string", diff --git a/augur/data/schema-translate-config.json b/augur/data/schema-translate-config.json index ab339e025..9a74490a7 100644 --- a/augur/data/schema-translate-config.json +++ b/augur/data/schema-translate-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/translate-config/v1" + }, "tree": { "description": "prebuilt Newick -- no tree will be built if provided", "type": "string", diff --git a/augur/data/schema-tree-config.json b/augur/data/schema-tree-config.json index c9a07e982..9a9ca547d 100644 --- a/augur/data/schema-tree-config.json +++ b/augur/data/schema-tree-config.json @@ -6,6 +6,9 @@ "type": "object", "additionalProperties": false, "properties": { + "$schema": { + "const": "https://nextstrain.org/schemas/augur/tree-config/v1" + }, "alignment": { "description": "alignment in fasta or VCF format", "type": "string", diff --git a/devel/regenerate-config-schemas b/devel/regenerate-config-schemas index e38c11704..c2979e105 100755 --- a/devel/regenerate-config-schemas +++ b/devel/regenerate-config-schemas @@ -56,14 +56,20 @@ def generate_schema_for_command(command_name: str, subparser: CustomArgumentPars properties = extract_properties(subparser) constraints = extract_constraints(subparser, properties) + schema_url = f'https://nextstrain.org/schemas/augur/{command_name.replace(" ", "-")}-config/{version}' schema: dict[str, Any] = { "_description": "This file is generated by devel/generate-config-schemas. Do not edit manually - edit the script instead.", "$schema": "http://json-schema.org/draft-07/schema#", - "$id": f'https://nextstrain.org/schemas/augur/{command_name.replace(" ", "-")}-config/{version}', + "$id": schema_url, "title": f"Configuration file to be supplied to `augur {command_name} {CONFIG_FILE_ARG}`", "type": "object", "additionalProperties": False, - "properties": properties, + "properties": { + "$schema": { + "const": schema_url, + }, + **properties, + }, } # NOTE: Nothing is required because everything can also be set in the CLI. # Though, the CLI help text still says required for options that can now be diff --git a/devel/regenerate-subsample-schema b/devel/regenerate-subsample-schema index 193f7d7da..c19b53d83 100755 --- a/devel/regenerate-subsample-schema +++ b/devel/regenerate-subsample-schema @@ -275,6 +275,9 @@ def create_schema(include_proximal: bool): # Define the schema structure. schema["properties"] = { + "$schema": { + "const": schema["$id"], + }, "defaults": {"$ref": "#/$defs/defaultProperties"}, "samples": { "type": "object", diff --git a/docs/_ext/config_schema_helpers.py b/docs/_ext/config_schema_helpers.py index 6e3336f27..ea31f8c8a 100644 --- a/docs/_ext/config_schema_helpers.py +++ b/docs/_ext/config_schema_helpers.py @@ -227,6 +227,10 @@ def run(self): tgroup.append(tbody) for prop_name in options.keys(): + # Keep the table focused on options from command line. + if prop_name == '$schema': + continue + prop_def = options[prop_name] row = nodes.row() tbody.append(row) From 8194a2065dd3256bba036f804c7e1b3e55737b11 Mon Sep 17 00:00:00 2001 From: Victor Lin Date: Fri, 11 Sep 2026 17:30:03 -0700 Subject: [PATCH 2/4] Reorganize schema helper functions Split the Augur-loading functionality of `load_json_schema()` into a higher level function `load_augur_json_schema()` so that the former can be kept lower level. There are now two external facing functions with clear intent: `load_augur_json_schema()` and `load_json_schema_locally()` --- augur/subsample.py | 6 +-- augur/util_support/node_data_file.py | 4 +- augur/validate.py | 60 +++++++++++++++++++--------- tests/test_validate.py | 21 +++++++--- 4 files changed, 62 insertions(+), 29 deletions(-) diff --git a/augur/subsample.py b/augur/subsample.py index 7bef66d4b..d41a4f849 100644 --- a/augur/subsample.py +++ b/augur/subsample.py @@ -25,7 +25,7 @@ from augur.io.metadata import DEFAULT_DELIMITERS, DEFAULT_ID_COLUMNS from augur.io.print import print_err from augur.utils import augur -from augur.validate import load_json_schema, validate_json, ValidateError +from augur.validate import load_augur_json_schema, validate_json, ValidateError from augur.io.print import print_debug BooleanFlags = Tuple[str, Optional[str]] @@ -181,7 +181,7 @@ def run(args: argparse.Namespace) -> None: worth it if a proper input reuse approach such as database/parquet file support is adopted: """ - schema_validator = load_json_schema("schema-subsample-config.json") + schema_validator = load_augur_json_schema("schema-subsample-config.json") config = _parse_config(args.config) try: @@ -497,7 +497,7 @@ def _get_sample_types(config: Dict[str, Any]) -> Dict[str, str]: Return a mapping of sample name to its schema variant (e.g. ``'filterSampleProperties'`` or ``'proximalSampleProperties'``). """ - schema = load_json_schema("schema-subsample-config.json").schema + schema = load_augur_json_schema("schema-subsample-config.json").schema variants = schema["properties"]["samples"]["patternProperties"]["^.+$"]["oneOf"] return { name: _best_matching_variant(variants, options, schema)[0] diff --git a/augur/util_support/node_data_file.py b/augur/util_support/node_data_file.py index 53fd0cc92..700384b40 100644 --- a/augur/util_support/node_data_file.py +++ b/augur/util_support/node_data_file.py @@ -6,7 +6,7 @@ from augur.io.file import open_file from augur.io.print import print_err from augur.types import ValidationMode -from augur.validate import validate_json, ValidateError, load_json_schema +from augur.validate import validate_json, ValidateError, load_augur_json_schema FILTERED_ATTRS = ["generated_by"] @@ -69,7 +69,7 @@ def validate(self): # validate_json will print any errors to STDERR validate_json( self.annotations, - load_json_schema("schema-annotations.json"), + load_augur_json_schema("schema-annotations.json"), self.fname, ) except ValidateError as err: diff --git a/augur/validate.py b/augur/validate.py index b60595799..202f769e4 100644 --- a/augur/validate.py +++ b/augur/validate.py @@ -42,12 +42,29 @@ def validation_failure(mode: ValidationMode): raise ValueError(f"unknown validation mode: {mode!r}") +def load_augur_json_schema(path: str): + """ + Load a JSON schema from the augur included set of schemas + (located in augur/data). + """ + with as_file(path) as file: + return load_json_schema_locally(file) + + def load_json_schema_locally(path): - # Schemas can reference other schemas via URL. Use local schema files for - # references to our own Augur schemas. This has advantages over URL: - # 1. it uses the schema expected by the current Augur version (URL may serve - # schema for a newer version) - # 2. it doesn't require an internet connection + """ + Load a JSON schema offline using local files. + + If the schema references Augur schemas via URL, the Augur schemas are served + by files bundled with the local Augur installation instead of what's served + at the remote URL. + + There are two advantages to this approach: + + 1. it uses the schema expected by the current Augur version (URL may serve + schema for a newer version) + 2. it doesn't require an internet connection + """ local_refs = { 'https://nextstrain.org/schemas/augur/annotations': "schema-annotations.json", 'https://nextstrain.org/schemas/dataset/root-sequence': "schema-export-root-sequence.json", @@ -60,17 +77,24 @@ def load_json_schema_locally(path): command_name = command.replace(' ', '-') local_refs[f"https://nextstrain.org/schemas/augur/{command_name}-config/{version}"] = f"schema-{command_name}-config.json" - return load_json_schema(path, refs=local_refs) + return _load_json_schema(path, refs=local_refs) -def load_json_schema(path, refs=None): - ''' - Load a JSON schema from the augur included set of schemas - (located in augur/data) - ''' +def _load_json_schema(path, refs=None): + """ + Load a schema from the path. + + Parameters + ---------- + path + Path to a JSON or YAML schema file. + refs : dict + Mapping of schema URI strings to local file paths (relative to + augur/data) used to resolve schema references offline. + """ is_yaml = Path(path).suffix in ('.yaml', '.yml') try: - with as_file(path) as file, open_file(file, "r") as fh: + with open_file(path, "r") as fh: if is_yaml: schema = YAML(typ="safe").load(fh) else: @@ -263,7 +287,7 @@ def grouped(iterable, key): def auspice_config_v2(config_json: Union[str,dict], **kwargs): - schema = load_json_schema("schema-auspice-config-v2.json") + schema = load_augur_json_schema("schema-auspice-config-v2.json") if isinstance(config_json, dict): config = config_json filename = "merged config" @@ -274,7 +298,7 @@ def auspice_config_v2(config_json: Union[str,dict], **kwargs): validate(config, schema, filename) def export_v2(main_json, **kwargs): - main_schema = load_json_schema_locally("schema-export-v2.json") + main_schema = load_augur_json_schema("schema-export-v2.json") if main_json.endswith("frequencies.json") or main_json.endswith("entropy.json") or main_json.endswith("sequences.json"): raise ValidateError("This validation subfunction is for the main `augur export v2` JSON only.") @@ -289,8 +313,8 @@ def export_v2(main_json, **kwargs): def export_v1(meta_json, tree_json, **kwargs): - meta_schema = load_json_schema("schema-export-v1-meta.json") - tree_schema = load_json_schema("schema-export-v1-tree.json") + meta_schema = load_augur_json_schema("schema-export-v1-meta.json") + tree_schema = load_augur_json_schema("schema-export-v1-tree.json") if not meta_json.endswith("_meta.json"): raise ValidateError("The metadata JSON pathname {} must end with '_meta.json'.".format(meta_json)) @@ -460,7 +484,7 @@ def validate_measurements_config(measurements): def measurements(measurements_json, **kwargs): - schema = load_json_schema("schema-measurements.json") + schema = load_augur_json_schema("schema-measurements.json") measurements = load_json(measurements_json) validate_json(measurements, schema, measurements_json) if not validate_measurements_config(measurements): @@ -469,7 +493,7 @@ def measurements(measurements_json, **kwargs): def measurements_collection_config(collection_config_json, **kwargs): - schema = load_json_schema("schema-measurements-collection-config.json") + schema = load_augur_json_schema("schema-measurements-collection-config.json") collection_config = load_json(collection_config_json) validate_json(collection_config, schema, collection_config_json) if not validate_collection_display_defaults(collection_config): diff --git a/tests/test_validate.py b/tests/test_validate.py index 9db8f8631..1e4a89e09 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -7,8 +7,9 @@ validate_collection_config_fields, validate_collection_display_defaults, validate_measurements_config, - load_json_schema, + load_augur_json_schema, load_json_schema_locally, + _load_json_schema, validate_json, ValidateError ) @@ -98,7 +99,7 @@ def test_validate_measurements_config_invalid_default_collection(self, example_m @pytest.fixture def genome_annotation_schema(): - return load_json_schema("schema-annotations.json") + return load_augur_json_schema("schema-annotations.json") class TestValidateGenomeAnnotations(): def test_negative_strand_nuc(self, capsys, genome_annotation_schema): @@ -197,8 +198,13 @@ def test_load_json_schema_locally(tmp_path): validate_json(valid_data, validator, "") -@pytest.mark.parametrize("ext", ["yaml", "yml"]) -def test_load_yaml_schema(tmp_path, ext): +def test_load_augur_json_schema(): + validator = load_augur_json_schema("schema-annotations.json") + assert validator.schema["$id"] == "https://nextstrain.org/schemas/augur/annotations" + + +@pytest.mark.parametrize("ext", ["json", "yaml", "yml"]) +def test_load_json_schema(tmp_path, ext): schema_file = tmp_path / f"test-schema.{ext}" raw_schema = { "type": "object", @@ -209,8 +215,11 @@ def test_load_yaml_schema(tmp_path, ext): } } with open(schema_file, "w") as f: - YAML(typ="safe").dump(raw_schema, f) + if ext == "json": + json.dump(raw_schema, f) + else: + YAML(typ="safe").dump(raw_schema, f) - validator = load_json_schema(schema_file) + validator = _load_json_schema(schema_file) valid_data = {"name": "test"} validate_json(valid_data, validator, "") From fdb33e6127480ec2445c3f64e60b6f6e3b1b3c1f Mon Sep 17 00:00:00 2001 From: Victor Lin Date: Fri, 11 Sep 2026 17:37:10 -0700 Subject: [PATCH 3/4] Load from schema URLs This will be useful for loading from URLs from a $schema key. --- augur/validate.py | 28 +++++++++++++++++++++++----- tests/test_validate.py | 6 ++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/augur/validate.py b/augur/validate.py index 202f769e4..21bc73674 100644 --- a/augur/validate.py +++ b/augur/validate.py @@ -42,13 +42,24 @@ def validation_failure(mode: ValidationMode): raise ValueError(f"unknown validation mode: {mode!r}") -def load_augur_json_schema(path: str): +def load_augur_json_schema(schema: str): """ Load a JSON schema from the augur included set of schemas - (located in augur/data). + (located in augur/data), specified by URL or local file name. """ - with as_file(path) as file: - return load_json_schema_locally(file) + local_refs = _get_local_refs() + if schema in local_refs: + # Convert URL to local reference + filename = local_refs[schema] + else: + # Assume everything else is already a local reference + filename = schema + + try: + with as_file(filename) as file: + return load_json_schema_locally(file) + except FileNotFoundError as err: + raise FileNotFoundError(f"No local file for schema: {schema}") from err def load_json_schema_locally(path): @@ -65,6 +76,13 @@ def load_json_schema_locally(path): schema for a newer version) 2. it doesn't require an internet connection """ + return _load_json_schema(path, refs=_get_local_refs()) + + +def _get_local_refs() -> dict[str, str]: + """ + Return a mapping of Augur schema URL to local file name under augur/data. + """ local_refs = { 'https://nextstrain.org/schemas/augur/annotations': "schema-annotations.json", 'https://nextstrain.org/schemas/dataset/root-sequence': "schema-export-root-sequence.json", @@ -77,7 +95,7 @@ def load_json_schema_locally(path): command_name = command.replace(' ', '-') local_refs[f"https://nextstrain.org/schemas/augur/{command_name}-config/{version}"] = f"schema-{command_name}-config.json" - return _load_json_schema(path, refs=local_refs) + return local_refs def _load_json_schema(path, refs=None): diff --git a/tests/test_validate.py b/tests/test_validate.py index 1e4a89e09..a78d515cf 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -202,6 +202,12 @@ def test_load_augur_json_schema(): validator = load_augur_json_schema("schema-annotations.json") assert validator.schema["$id"] == "https://nextstrain.org/schemas/augur/annotations" + validator = load_augur_json_schema("https://nextstrain.org/schemas/augur/tree-config/v1") + assert validator.schema["$id"] == "https://nextstrain.org/schemas/augur/tree-config/v1" + + with pytest.raises(FileNotFoundError, match="No local file for schema"): + load_augur_json_schema("https://nextstrain.org/schemas/augur/invalid-config/v1") + @pytest.mark.parametrize("ext", ["json", "yaml", "yml"]) def test_load_json_schema(tmp_path, ext): From bb28bc598e943dc13b26d5a527ac6ecc1c0fec81 Mon Sep 17 00:00:00 2001 From: Victor Lin Date: Tue, 15 Sep 2026 16:18:33 -0700 Subject: [PATCH 4/4] Add generic get_referenced_files() Now that other commands accept config files with filepaths marked in schema, the generic get_referenced_files() function is necessary for adoption in Snakemake workflows where referenced files are to be listed as inputs for file change detection. Some changes were necessary as the previous subsample-specific function hardcodes the schema within the function. The generic function solves this by requiring the $schema key in the config. Keeping the subsample-specific function around because it's used by some workflows. Will remove once those workflows are updated to use the generic function. --- augur/config.py | 198 +++++++++++++++++++++++++++++++++++++++++ augur/subsample.py | 204 +++---------------------------------------- augur/validate.py | 4 +- tests/test_config.py | 76 ++++++++++++++++ 4 files changed, 289 insertions(+), 193 deletions(-) diff --git a/augur/config.py b/augur/config.py index b132cebaf..54cf02bae 100644 --- a/augur/config.py +++ b/augur/config.py @@ -1,10 +1,14 @@ """ Helpers for YAML-based configuration files. """ +import os from pathlib import Path +from ruamel.yaml import YAML, YAMLError, constructor from textwrap import dedent +from typing import Any, Optional from augur.errors import AugurError from augur.io.print import indented_list +from augur.validate import load_augur_json_schema, ValidateError COMMAND_SCHEMAS = { "align": "v1", @@ -17,6 +21,200 @@ "tree": "v1", } + +def get_referenced_files(config_file: str | Path) -> set[str]: + """Get the files referenced in an augur config file. + + Extracts and resolves all filepath values referenced in the config. + + Parameters + ---------- + config_file + Path to the config file. + + Returns + ------- + set + Resolved filepaths + """ + config = parse_config(config_file) + + schema_ref = config.get("$schema") + if not schema_ref: + raise AugurError(f"The configuration file {str(config_file)!r} does not specify a '$schema'.") + + try: + schema_validator = load_augur_json_schema(str(schema_ref)) + except (FileNotFoundError, ValidateError) as e: + raise AugurError(f"Schema {schema_ref!r} not found: {e}") from e + + # Resolve filepaths. + search_paths = get_search_paths(config_file) + config, filepaths = resolve_filepaths(config, search_paths, schema_validator.schema) + + return set(filepaths) + + +def parse_config(filename: str | Path) -> dict[str, Any]: + # Create a custom YAML constructor to treat timestamps as strings. + class CustomConstructor(constructor.SafeConstructor): + pass + def string_constructor(loader, node): + return loader.construct_scalar(node) + CustomConstructor.add_constructor('tag:yaml.org,2002:timestamp', string_constructor) + + yaml = YAML(typ="safe") + yaml.Constructor = CustomConstructor + + with open(filename) as f: + try: + config = yaml.load(f) + except YAMLError as e: + raise AugurError(f"The configuration file {filename!r} is not valid YAML.\n" + str(e)) from e + + return config + + +def get_search_paths(config_file: str | Path) -> list[Path]: + """ + Returns the paths to search for relative filepaths in config. + """ + default = [ + Path(config_file).parent, + Path.cwd(), + ] + + from_env = os.environ.get('AUGUR_SEARCH_PATHS') + + if from_env: + return [ + *(Path(p) for p in from_env.split(':')), + *default, + ] + + return default + + +def resolve_filepaths( + config: dict[str, Any], + search_paths: list[Path], + schema: dict[str, Any], + root_schema: Optional[dict[str, Any]] = None, +) -> tuple[dict[str, Any], list[str]]: + """ + Resolve filepaths in config. + + Recursively walks the config alongside the schema to determine which fields + contain filepaths, resolves them, and collects the resolved filepaths. + """ + if root_schema is None: + root_schema = schema + + filepaths = [] + + # Get properties schema for current section + properties = schema.get("properties", {}) + pattern_properties = schema.get("patternProperties", {}) + + for key, value in config.items(): + if key == "$schema" or key.startswith("_"): + continue + prop_schema = properties.get(key) + + if not prop_schema and pattern_properties: + # Use first pattern property schema (for dynamic keys like samples) + prop_schema = next(iter(pattern_properties.values())) + + # Get referenced property schema + if ref := prop_schema.get("$ref"): + prop_schema = _get_referenced_schema(ref, root_schema) + elif "oneOf" in prop_schema and isinstance(value, dict): + _, prop_schema = best_matching_variant(prop_schema["oneOf"], value, root_schema) + + # Resolve filepath + if _is_filepath(prop_schema): + if isinstance(value, list): + config[key] = [str(resolve_filepath(Path(v), search_paths)) for v in value] + filepaths.extend(config[key]) + elif isinstance(value, str): + config[key] = str(resolve_filepath(Path(value), search_paths)) + filepaths.append(config[key]) + + # Recurse into config section + elif isinstance(value, dict): + config[key], downstream_filepaths = resolve_filepaths( + value, search_paths, prop_schema, root_schema + ) + filepaths.extend(downstream_filepaths) + + return config, filepaths + + +def best_matching_variant( + variants: list[dict[str, Any]], + value: dict[str, Any], + root_schema: dict[str, Any], +) -> tuple[str, dict[str, Any]]: + """ + Given a oneOf list of schema variants (typically $ref entries), resolve each + and return the one whose properties best match the keys in *value*, as a + (name, schema) tuple. + + The name is the last component of the matched ``$ref`` (e.g. + ``"filterSampleProperties"``). + """ + value_keys = set(value.keys()) + best_schema = None + best_ref = None + best_overlap = -1 + for variant in variants: + ref = variant.get("$ref") + if ref: + resolved = _get_referenced_schema(ref, root_schema) + else: + resolved = variant + props = set(resolved.get("properties", {}).keys()) + overlap = len(value_keys & props) + if overlap > best_overlap: + best_overlap = overlap + best_schema = resolved + best_ref = ref + if not best_schema: + raise AugurError("Couldn't match oneOf schema for config dict") + return (best_ref.rsplit("/", 1)[-1], best_schema) + + +def _get_referenced_schema( + ref: str, + root_schema: dict[str, Any], +) -> dict[str, Any]: + """ + Resolve a JSON schema reference. Example: '#/$defs/filterSampleProperties' + """ + keys = ref.lstrip("#/").split("/") + schema = root_schema + for key in keys: + schema = schema[key] + return schema + + +def _is_filepath(prop_schema: dict[str, Any]) -> bool: + """ + Check if the property schema declares it is a filepath. + """ + # Direct 'format: filepath' + if prop_schema.get("format") == "filepath": + return True + + # Check oneOf variants for 'format: filepath' + if "oneOf" in prop_schema: + for variant in prop_schema["oneOf"]: + if variant.get("format") == "filepath": + return True + + return False + + def resolve_filepath( path: Path, search_paths: list[Path], diff --git a/augur/subsample.py b/augur/subsample.py index d41a4f849..195148baf 100644 --- a/augur/subsample.py +++ b/augur/subsample.py @@ -12,15 +12,18 @@ import subprocess import sys import tempfile -from ruamel.yaml import YAML, YAMLError, constructor from collections import defaultdict, deque from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, Future, wait -from pathlib import Path from textwrap import dedent -from typing import Any, Dict, List, Optional, Set, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from augur import filter as augur_filter from augur.argparse_ import ExtendOverwriteDefault, SKIP_AUTO_DEFAULT_IN_HELP -from augur.config import resolve_filepath +from augur.config import ( + best_matching_variant, + get_search_paths, + parse_config, + resolve_filepaths, +) from augur.errors import AugurError from augur.io.metadata import DEFAULT_DELIMITERS, DEFAULT_ID_COLUMNS from augur.io.print import print_err @@ -182,7 +185,7 @@ def run(args: argparse.Namespace) -> None: support is adopted: """ schema_validator = load_augur_json_schema("schema-subsample-config.json") - config = _parse_config(args.config) + config = parse_config(args.config) try: validate_json(config, schema_validator, args.config) @@ -198,8 +201,8 @@ def run(args: argparse.Namespace) -> None: raise AugurError("Proximal sampling for AA sequences is not yet supported.") # Resolve filepaths. - search_paths = _get_search_paths(args.config) - config, filepaths = _resolve_filepaths(config, search_paths, schema_validator.schema) + search_paths = get_search_paths(args.config) + config, filepaths = resolve_filepaths(config, search_paths, schema_validator.schema) print_debug(f"\nResolved filepaths: {filepaths}") # Construct sample objects for augur filter calls @@ -286,31 +289,6 @@ def run(args: argparse.Namespace) -> None: s.remove_temporary_files() -def get_referenced_files(config_file: str) -> Set[str]: - """Get the files referenced in a subsample config file. - - Extracts and resolves all filepath values referenced in the config, - including defaults and individual sample options. - - Parameters - ---------- - config_file - Path to the subsample config file. - - Returns - ------- - set - Resolved filepaths - """ - schema_validator = load_json_schema("schema-subsample-config.json") - config = _parse_config(config_file) - - # Resolve filepaths. - search_paths = _get_search_paths(config_file) - config, filepaths = _resolve_filepaths(config, search_paths, schema_validator.schema) - - return set(filepaths) - def requires_aligned_sequences(config_file: str) -> bool: """NOTE: This function may change without warning & is for internal Snakemake development / testing purposes. @@ -328,170 +306,12 @@ def requires_aligned_sequences(config_file: str) -> bool: bool Does augur subsample require aligned sequences? """ - config = _parse_config(config_file) + config = parse_config(config_file) return _includes_proximal_sample(config) -def _parse_config(filename: str) -> Dict[str, Any]: - # Create a custom YAML constructor to treat timestamps as strings. - class CustomConstructor(constructor.SafeConstructor): - pass - def string_constructor(loader, node): - return loader.construct_scalar(node) - CustomConstructor.add_constructor('tag:yaml.org,2002:timestamp', string_constructor) - - yaml = YAML(typ="safe") - yaml.Constructor = CustomConstructor - - with open(filename) as f: - try: - config = yaml.load(f) - except YAMLError as e: - raise AugurError(f"The configuration file {filename!r} is not valid YAML.\n" + str(e)) from e - - return config - -def _get_search_paths( - config_file: str, -) -> List[Path]: - """ - Returns the paths to search for relative filepaths in config. - """ - default = [ - Path(config_file).parent, - Path.cwd(), - ] - - from_env = os.environ.get('AUGUR_SEARCH_PATHS') - - if from_env: - return [ - *(Path(p) for p in from_env.split(':')), - *default, - ] - - return default - - -def _resolve_filepaths( - config: Dict[str, Any], - search_paths: List[Path], - schema: Dict[str, Any], - root_schema: Optional[Dict[str, Any]] = None, -) -> Tuple[Dict[str, Any], List[str]]: - """ - Resolve filepaths in config. - - Recursively walks the config alongside the schema to determine which fields - contain filepaths, resolves them, and collects the resolved filepaths. - """ - if root_schema is None: - root_schema = schema - - filepaths = [] - - # Get properties schema for current section - properties = schema.get("properties", {}) - pattern_properties = schema.get("patternProperties", {}) - - for key, value in config.items(): - if key.startswith("_"): - continue - prop_schema = properties.get(key) - - if not prop_schema and pattern_properties: - # Use first pattern property schema (for dynamic keys like samples) - prop_schema = next(iter(pattern_properties.values())) - - # Get referenced property schema - if ref := prop_schema.get("$ref"): - prop_schema = _get_referenced_schema(ref, root_schema) - elif "oneOf" in prop_schema and isinstance(value, dict): - _, prop_schema = _best_matching_variant(prop_schema["oneOf"], value, root_schema) - - # Resolve filepath - if _is_filepath(prop_schema): - if isinstance(value, list): - config[key] = [str(resolve_filepath(Path(v), search_paths)) for v in value] - filepaths.extend(config[key]) - elif isinstance(value, str): - config[key] = str(resolve_filepath(Path(value), search_paths)) - filepaths.append(config[key]) - - # Recurse into config section - elif isinstance(value, dict): - config[key], downstream_filepaths = _resolve_filepaths( - value, search_paths, prop_schema, root_schema - ) - filepaths.extend(downstream_filepaths) - - return config, filepaths - def _includes_proximal_sample(config: Dict[str, Any]) -> bool: return 'proximalSampleProperties' in _get_sample_types(config).values() -def _get_referenced_schema( - ref: str, - root_schema: Dict[str, Any], -) -> Dict[str, Any]: - """ - Resolve a JSON schema reference. Example: '#/$defs/filterSampleProperties' - """ - keys = ref.lstrip("#/").split("/") - schema = root_schema - for key in keys: - schema = schema[key] - return schema - - -def _is_filepath(prop_schema: Dict[str, Any]) -> bool: - """ - Check if the property schema declares it is a filepath. - """ - # Direct 'format: filepath' - if prop_schema.get("format") == "filepath": - return True - - # Check oneOf variants for 'format: filepath' - if "oneOf" in prop_schema: - for variant in prop_schema["oneOf"]: - if variant.get("format") == "filepath": - return True - - return False - -def _best_matching_variant( - variants: List[Dict[str, Any]], - value: Dict[str, Any], - root_schema: Dict[str, Any], -) -> Tuple[str, Dict[str, Any]]: - """ - Given a oneOf list of schema variants (typically $ref entries), resolve each - and return the one whose properties best match the keys in *value*, as a - (name, schema) tuple. - - The name is the last component of the matched ``$ref`` (e.g. - ``"filterSampleProperties"``). - """ - value_keys = set(value.keys()) - best_schema = None - best_ref = None - best_overlap = -1 - for variant in variants: - ref = variant.get("$ref") - if ref: - resolved = _get_referenced_schema(ref, root_schema) - else: - resolved = variant - props = set(resolved.get("properties", {}).keys()) - overlap = len(value_keys & props) - if overlap > best_overlap: - best_overlap = overlap - best_schema = resolved - best_ref = ref - if not best_schema: - raise AugurError("Couldn't match oneOf schema for config dict") - return (best_ref.rsplit("/", 1)[-1], best_schema) - def _get_sample_types(config: Dict[str, Any]) -> Dict[str, str]: """ Return a mapping of sample name to its schema variant (e.g. @@ -500,7 +320,7 @@ def _get_sample_types(config: Dict[str, Any]) -> Dict[str, str]: schema = load_augur_json_schema("schema-subsample-config.json").schema variants = schema["properties"]["samples"]["patternProperties"]["^.+$"]["oneOf"] return { - name: _best_matching_variant(variants, options, schema)[0] + name: best_matching_variant(variants, options, schema)[0] for name, options in config.get("samples", {}).items() } diff --git a/augur/validate.py b/augur/validate.py index 21bc73674..a02fe7c79 100644 --- a/augur/validate.py +++ b/augur/validate.py @@ -15,7 +15,6 @@ from textwrap import indent from typing import Iterable, Union from augur.argparse_ import add_subparser -from augur.config import COMMAND_SCHEMAS from augur.data import as_file from augur.io.file import open_file from augur.io.print import print_err @@ -83,6 +82,9 @@ def _get_local_refs() -> dict[str, str]: """ Return a mapping of Augur schema URL to local file name under augur/data. """ + # Import here to avoid circular top-level imports. + from augur.config import COMMAND_SCHEMAS + local_refs = { 'https://nextstrain.org/schemas/augur/annotations': "schema-annotations.json", 'https://nextstrain.org/schemas/dataset/root-sequence': "schema-export-root-sequence.json", diff --git a/tests/test_config.py b/tests/test_config.py index 7c78c9ca5..e602ef622 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,6 +3,8 @@ from ruamel.yaml import YAML from textwrap import dedent from augur import make_parser +from augur.config import get_referenced_files +from augur.errors import AugurError def write_config_file(tmp_path: Path, config: dict) -> Path: @@ -174,3 +176,77 @@ def test_config_error_with_duplicate(tmp_path, capsys): assert exc_info.value.code == 2 captured = capsys.readouterr() assert 'found duplicate key "timetree"' in captured.err + + +def test_get_referenced_files_tree(tmp_path): + """ + Test get_referenced_files with augur tree config. + """ + alignment = tmp_path / "alignment.fasta" + alignment.touch() + config_file = write_config_file(tmp_path, { + "$schema": "https://nextstrain.org/schemas/augur/tree-config/v1", + "alignment": "alignment.fasta", + }) + + assert get_referenced_files(config_file) == {str(alignment.resolve())} + + +def test_get_referenced_files_subsample(tmp_path): + """ + Test get_referenced_files with nested structure in subsample config. + """ + include_file = tmp_path / "include.txt" + include_file.touch() + exclude_file = tmp_path / "exclude.txt" + exclude_file.touch() + + config_file = write_config_file(tmp_path, { + "$schema": "https://nextstrain.org/schemas/augur/subsample-config/v1", + "defaults": { + "exclude": "exclude.txt", + }, + "samples": { + "sample_a": { + "include": ["include.txt"], + }, + }, + }) + + assert get_referenced_files(config_file) == {str(include_file.resolve()), str(exclude_file.resolve())} + + +def test_get_referenced_files_missing_schema(tmp_path): + """ + Test that an error is raised when $schema is missing. + """ + config_file = write_config_file(tmp_path, {"alignment": "alignment.fasta"}) + + with pytest.raises(AugurError, match=r"does not specify a '\$schema'"): + get_referenced_files(config_file) + + +def test_get_referenced_files_nonexistent_command(tmp_path): + """ + Test that an error is raised when the inferred command schema doesn't exist. + """ + config_file = write_config_file(tmp_path, { + "$schema": "https://nextstrain.org/schemas/augur/nonexistent-config/v1", + "alignment": "alignment.fasta", + }) + + with pytest.raises(AugurError, match="not found"): + get_referenced_files(config_file) + + +def test_get_referenced_files_missing_file(tmp_path): + """ + Test that an error is raised when a referenced file cannot be resolved. + """ + config_file = write_config_file(tmp_path, { + "$schema": "https://nextstrain.org/schemas/augur/tree-config/v1", + "alignment": "nonexistent.fasta", + }) + + with pytest.raises(AugurError, match="not resolvable from any of the following paths"): + get_referenced_files(config_file)