Merge pull request #126 from crim-ca/simplify-populator-discovery

Simplify populator discovery

Author: mishaschwartz

CHANGES.md

@@ -2,7 +2,14 @@
 
 ## [Unreleased](https://github.com/crim-ca/stac-populator) (latest)
 
-<!-- insert list items of new changes here -->
+* Simplify populator implementation discovery and reduce boilerplate for new implementations.
+* Add a classmethod to add additional CLI args to the `STACpopulatorBase` class (replaces `add_parser_args`)
+* Add a classmethod to run the populator based on CLI args to the `STACpopulatorBase` class (replaces `runner`)
+* Add shared implementations for these two classmethods to a `THREDDSpopulatorBase` class that can be
+  used by all populators that pull data from THREDDS
+* Rename `populator_base.py` to `populators.py` since it now contains more than just the base populator
+* No longer require new implementations to list populator modules in `implementations.__init__.py`
+* Implementation discovery is now done by keeping track of concrete implementations of the STACpopulatorBase class
 
 ## [0.12.0](https://github.com/crim-ca/stac-populator/tree/0.12.0) (2025-11-20)
 

STACpopulator/cli.py

@@ -3,26 +3,22 @@
 import importlib
 import sys
 import warnings
-from types import ModuleType
+from pathlib import Path
 from typing import get_args
 
 import pystac
 import requests
 
-from STACpopulator import __version__, implementations
-from STACpopulator.collection_update import UpdateModes, UpdateModesOptional, update_api_collection
+import STACpopulator.implementations
+from STACpopulator import __version__
+from STACpopulator.collection_update import UpdateModes, update_api_collection
 from STACpopulator.exceptions import STACPopulatorError
 from STACpopulator.export import export_catalog
 from STACpopulator.log import add_logging_options, setup_logging
+from STACpopulator.populators import STACpopulatorBase
 from STACpopulator.request_utils import add_request_options, apply_request_options
 
 
-def _extra_parser_argument(arg: str) -> tuple[str, str]:
-    if "=" in arg:
-        return tuple(a.strip() for a in arg.split("=", 1))
-    raise argparse.ArgumentTypeError("--extra-parser-arguments must be in the form 'key=value'")
-
-
 def add_parser_args(parser: argparse.ArgumentParser) -> None:
     """Add parser arguments to the argument parser."""
     parser.add_argument(
@@ -48,52 +44,10 @@ def add_parser_args(parser: argparse.ArgumentParser) -> None:
     populators_subparser = run_parser.add_subparsers(
         title="populator", dest="populator", description="Implementation to run."
     )
-    for implementation_module_name, module in implementation_modules().items():
-        implementation_parser = populators_subparser.add_parser(implementation_module_name)
-        module.add_parser_args(implementation_parser)
-        implementation_parser.add_argument(
-            "-x",
-            "--extra-item-parsers",
-            action="append",
-            help="Functions that may modify items before upload. "
-            "Should be specified in the form 'module:function_name' "
-            "and have the signature function(item: dict, **kw)",
-        )
-        implementation_parser.add_argument(
-            "-X",
-            "--extra-collection-parsers",
-            action="append",
-            help="Functions that may modify collections before upload. "
-            "Should be specified in the form 'module:function_name' or "
-            "path/to/python/file.py:function_name. Functions should "
-            "have the signature function(collection: dict, **kw) -> None "
-            "and should modify the collection dict in place.",
-        )
-        implementation_parser.add_argument(
-            "-a",
-            "--extra-parser-arguments",
-            action="append",
-            type=_extra_parser_argument,
-            help="Extra keyword arguments that should be passed to extra "
-            "item and collection function as "
-            "keyword arguments. "
-            "Should be specified in the form 'key=value'",
-        )
-        implementation_parser.add_argument(
-            "--update-collection-mode",
-            dest="update_collection",
-            choices=get_args(UpdateModesOptional),
-            default="none",
-            help="Update collection information based on new items created or updated by this populator. "
-            "Only applies if --update is also set.",
-        )
-        implementation_parser.add_argument(
-            "--exclude-summary",
-            nargs="*",
-            action="extend",
-            default=[],
-            help="Exclude these properties when updating collection summaries. ",
-        )
+    for name, populator in populators().items():
+        implementation_parser = populators_subparser.add_parser(name)
+        implementation_parser.description = getattr(populator, "description", name)
+        populator.update_parser_args(implementation_parser)
     update_parser = commands_subparser.add_parser(
         "update-collection", description="Update collection information based on items in the collection"
     )
@@ -123,21 +77,24 @@ def add_parser_args(parser: argparse.ArgumentParser) -> None:
 
 
 @functools.cache
-def implementation_modules() -> dict[str, ModuleType]:
+def populators() -> dict[str, STACpopulatorBase]:
     """
     Try to load implementations.
 
     If one fails (i.e. due to missing dependencies) continue loading others.
     """
-    modules = {}
-    for implementation_module_name in implementations.__all__:
-        try:
-            modules[implementation_module_name] = importlib.import_module(
-                f".{implementation_module_name}", implementations.__package__
-            )
-        except STACPopulatorError as e:
-            warnings.warn(f"Could not load extension {implementation_module_name} because of error {e}")
-    return modules
+    impl_path = Path(STACpopulator.implementations.__path__[0])
+    for path in impl_path.glob("**/*.py"):
+        if path.name == "__init__.py":
+            path = path.parent
+        rel_path = path.relative_to(impl_path)
+        if str(rel_path) != ".":
+            module_path = str(rel_path.with_suffix("")).replace("/", ".").replace("-", "_")
+            try:
+                importlib.import_module(f"STACpopulator.implementations.{module_path}")
+            except STACPopulatorError as e:
+                warnings.warn(f"Could not load extension {rel_path} because of error {e}")
+    return {getattr(klass, "name", klass.__name__): klass for klass in STACpopulatorBase.concrete_subclasses()}
 
 
 def run(ns: argparse.Namespace) -> int:
@@ -148,7 +105,7 @@ def run(ns: argparse.Namespace) -> int:
         if ns.command == "run":
             if ns.stac_version:
                 pystac.set_stac_version(ns.stac_version)
-            return implementation_modules()[ns.populator].runner(ns, session) or 0
+            return populators()[ns.populator].run(ns, session) or 0
         elif ns.command == "update_collection":
             return update_api_collection(ns.mode, ns.stac_collection_uri, ns.exclude_summary) or 0
         else:

STACpopulator/implementations/CMIP6_UofT/__init__.py

@@ -1,3 +0,0 @@
-from .add_CMIP6 import add_parser_args, runner
-
-__all__ = ["add_parser_args", "runner"]

STACpopulator/implementations/CMIP6_UofT/add_CMIP6.py

@@ -1,25 +1,25 @@
-import argparse
 import json
 import logging
 from typing import Any, MutableMapping, Union
 
 from pystac import STACValidationError
 from pystac.extensions.datacube import DatacubeExtension
-from requests.sessions import Session
 
 from STACpopulator.extensions.cmip6 import CMIP6Helper, CMIP6Properties
 from STACpopulator.extensions.datacube import DataCubeHelper
 from STACpopulator.extensions.thredds import THREDDSExtension, THREDDSHelper
-from STACpopulator.input import ErrorLoader, THREDDSLoader
 from STACpopulator.models import GeoJSONPolygon
-from STACpopulator.populator_base import STACpopulatorBase
+from STACpopulator.populators import THREDDSPopulator
 
 LOGGER = logging.getLogger(__name__)
 
 
-class CMIP6populator(STACpopulatorBase):
+class CMIP6populator(THREDDSPopulator):
     """Populator that creates STAC objects representing CMIP6 data from a THREDDS catalog."""
 
+    name = "CMIP6_UofT"
+    description = "CMIP6 STAC populator from a THREDDS catalog or NCML XML."
+
     item_properties_model = CMIP6Properties
     item_geometry_model = GeoJSONPolygon
 
@@ -64,51 +64,3 @@ def create_stac_item(
 
         # print(json.dumps(item.to_dict()))
         return json.loads(json.dumps(item.to_dict()))
-
-
-def add_parser_args(parser: argparse.ArgumentParser) -> None:
-    """Add additional CLI arguments to the argument parser."""
-    parser.description = "CMIP6 STAC populator from a THREDDS catalog or NCML XML."
-    parser.add_argument("stac_host", help="STAC API URL")
-    parser.add_argument("href", help="URL to a THREDDS catalog or a NCML XML with CMIP6 metadata.")
-    parser.add_argument("--update", action="store_true", help="Update collection and its items")
-    parser.add_argument(
-        "--mode",
-        choices=["full", "single"],
-        default="full",
-        help="Operation mode, processing the full dataset or only the single reference.",
-    )
-    parser.add_argument(
-        "--config",
-        type=str,
-        help=(
-            "Override configuration file for the populator. "
-            "By default, uses the adjacent configuration to the implementation class."
-        ),
-    )
-
-
-def runner(ns: argparse.Namespace, session: Session) -> int:
-    """Run the populator."""
-    LOGGER.info(f"Arguments to call: {vars(ns)}")
-
-    if ns.mode == "full":
-        data_loader = THREDDSLoader(ns.href, session=session)
-    else:
-        # To be implemented
-        data_loader = ErrorLoader()
-
-    c = CMIP6populator(
-        ns.stac_host,
-        data_loader,
-        update=ns.update,
-        session=session,
-        config_file=ns.config,
-        extra_item_parsers=ns.extra_item_parsers,
-        extra_collection_parsers=ns.extra_collection_parsers,
-        extra_parser_arguments=ns.extra_parser_arguments,
-        update_collection=ns.update_collection,
-        exclude_summaries=ns.exclude_summary,
-    )
-    c.ingest()
-    return 0

STACpopulator/implementations/CORDEXCMIP6_Ouranos/__init__.py

@@ -1,3 +0,0 @@
-from STACpopulator.implementations.CORDEXCMIP6_Ouranos.add_CORDEX6 import add_parser_args, runner
-
-__all__ = ["add_parser_args", "runner"]

STACpopulator/implementations/CORDEXCMIP6_Ouranos/add_CORDEX6.py

@@ -1,69 +1,22 @@
-import argparse
 import logging
 from typing import Any
 
-from requests.sessions import Session
-
 from STACpopulator.extensions.cordex6 import Cordex6DataModel
-from STACpopulator.input import ErrorLoader, THREDDSLoader
-from STACpopulator.populator_base import STACpopulatorBase
+from STACpopulator.populators import THREDDSPopulator
 
 LOGGER = logging.getLogger(__name__)
 
 
-class CORDEX_STAC_Populator(STACpopulatorBase):
+class CORDEX_STAC_Populator(THREDDSPopulator):
     """Populator that creates STAC objects representing CORDEX data from a THREDDS catalog."""
 
+    name = "CORDEXCMIP6_Ouranos"
+    description = "CMIP6-CORDEX STAC populator from a THREDDS catalog or NCML XML."
+
     data_model = Cordex6DataModel
     item_geometry_model = None  # Unnecessary, but kept for consistency
 
     def create_stac_item(self, item_name: str, item_data: dict[str, Any]) -> dict[str, Any]:
         """Return a STAC item."""
         dm = self.data_model.from_data(item_data)
         return dm.stac_item()
-
-
-# TODO: This probably doesn't need to be copied for every implementation, right ?
-def add_parser_args(parser: argparse.ArgumentParser) -> None:
-    """Add additional CLI arguments to the argument parser."""
-    parser.description = "CMIP6-CORDEX STAC populator from a THREDDS catalog or NCML XML."
-    parser.add_argument("stac_host", help="STAC API URL")
-    parser.add_argument("href", help="URL to a THREDDS catalog or a NCML XML with CMIP6 metadata.")
-    parser.add_argument("--update", action="store_true", help="Update collection and its items")
-    parser.add_argument(
-        "--mode",
-        choices=["full", "single"],
-        default="full",
-        help="Operation mode, processing the full dataset or only the single reference.",
-    )
-    parser.add_argument(
-        "--config",
-        type=str,
-        help=(
-            "Override configuration file for the populator. "
-            "By default, uses the adjacent configuration to the implementation class."
-        ),
-    )
-
-
-def runner(ns: argparse.Namespace, session: Session) -> int:
-    """Run the populator."""
-    LOGGER.info(f"Arguments to call: {vars(ns)}")
-
-    if ns.mode == "full":
-        data_loader = THREDDSLoader(ns.href, session=session)
-    else:
-        # To be implemented
-        data_loader = ErrorLoader()
-
-    c = CORDEX_STAC_Populator(
-        ns.stac_host,
-        data_loader,
-        update=ns.update,
-        session=session,
-        config_file=ns.config,
-        update_collection=ns.update_collection,
-        exclude_summaries=ns.exclude_summary,
-    )
-    c.ingest()
-    return 0

STACpopulator/implementations/DirectoryLoader/__init__.py

@@ -1,3 +0,0 @@
-from STACpopulator.implementations.DirectoryLoader.crawl_directory import add_parser_args, runner
-
-__all__ = ["add_parser_args", "runner"]