Support generic invocation of package_discovery functions

colcon_core/package_selection/__init__.py

@@ -71,18 +71,26 @@ def select_packages(self, *, args, decorators):
         raise NotImplementedError()
 
 
-def add_arguments(parser):
+def add_arguments(
+    parser, *, discovery_extensions=None, selection_extensions=None,
+):
     """
     Add the command line arguments for the package selection extensions.
 
     The function will call :function:`add_package_discovery_arguments` to add
     the package discovery arguments.
 
     :param parser: The argument parser
+    :param discovery_extensions: The package discovery extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_discovery_extensions`
+    :param selection_extensions: The package selection extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_selection_extensions`
     """
-    add_package_discovery_arguments(parser)
+    add_package_discovery_arguments(parser, extensions=discovery_extensions)
 
-    _add_package_selection_arguments(parser)
+    _add_package_selection_arguments(parser, extensions=selection_extensions)
 
 
 def get_package_selection_extensions(*, group_name=None):
@@ -101,15 +109,16 @@ def get_package_selection_extensions(*, group_name=None):
     return order_extensions_by_priority(extensions)
 
 
-def _add_package_selection_arguments(parser):
+def _add_package_selection_arguments(parser, *, extensions=None):
     """
     Add the command line arguments for the package selection extensions.
 
     :param parser: The argument parser
     """
-    package_selection_extensions = get_package_selection_extensions()
+    if extensions is None:
+        extensions = get_package_selection_extensions()
     group = parser.add_argument_group(title='Package selection arguments')
-    for extension in package_selection_extensions.values():
+    for extension in extensions.values():
         try:
             retval = extension.add_arguments(parser=group)
             assert retval is None, 'add_arguments() should return None'
@@ -125,7 +134,9 @@ def _add_package_selection_arguments(parser):
 def get_packages(
     args, *,
     additional_argument_names=None,
-    direct_categories=None, recursive_categories=None
+    direct_categories=None, recursive_categories=None,
+    discovery_extensions=None, identification_extensions=None,
+    augmentation_extensions=None, selection_extensions=None,
 ):
     """
     Get the selected package decorators in topological order.
@@ -141,17 +152,35 @@ def get_packages(
     :param Iterable[str]|Mapping[str, Iterable[str]] recursive_categories:
       The names of the recursive categories, optionally mapped from the
       immediate upstream category which included the dependency
+    :param discovery_extensions: The package discovery extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_discovery_extensions`
+    :param identification_extensions: The package identification extensions to
+      use, if `None` is passed use the extensions provided by
+      :function:`get_package_identification_extensions`
+    :param augmentation_extensions: The package augmentation extensions, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_augmentation_extensions`
+    :param selection_extensions: The package selection extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_selection_extensions`
     :rtype: list
     :raises RuntimeError: if the returned set of packages contains duplicates
       package names
     """
     descriptors = get_package_descriptors(
-        args, additional_argument_names=additional_argument_names)
+        args, additional_argument_names=additional_argument_names,
+        discovery_extensions=discovery_extensions,
+        identification_extensions=identification_extensions,
+        augmentation_extensions=augmentation_extensions,
+        selection_extensions=selection_extensions)
     decorators = topological_order_packages(
         descriptors,
         direct_categories=direct_categories,
         recursive_categories=recursive_categories)
-    select_package_decorators(args, decorators)
+    select_package_decorators(
+        args, decorators,
+        selection_extensions=selection_extensions)
 
     # check for duplicate package names
     pkgs = [m.descriptor for m in decorators if m.selected]
@@ -169,7 +198,11 @@ def get_packages(
     return decorators
 
 
-def get_package_descriptors(args, *, additional_argument_names=None):
+def get_package_descriptors(
+    args, *, additional_argument_names=None, discovery_extensions=None,
+    identification_extensions=None, augmentation_extensions=None,
+    selection_extensions=None,
+):
     """
     Get the package descriptors.
 
@@ -181,24 +214,44 @@ def get_package_descriptors(args, *, additional_argument_names=None):
 
     :param additional_argument_names: A list of additional arguments to
       consider
+    :param discovery_extensions: The package discovery extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_discovery_extensions`
+    :param identification_extensions: The package identification extensions to
+      use, if `None` is passed use the extensions provided by
+      :function:`get_package_identification_extensions`
+    :param augmentation_extensions: The package augmentation extensions, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_augmentation_extensions`
+    :param selection_extensions: The package selection extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_selection_extensions`
     :returns: set of
       :py:class:`colcon_core.package_descriptor.PackageDescriptor`
     :rtype: set
     """
-    extensions = get_package_identification_extensions()
-    descriptors = discover_packages(args, extensions)
+    if identification_extensions is None:
+        identification_extensions = get_package_identification_extensions()
+    descriptors = discover_packages(
+        args, identification_extensions,
+        discovery_extensions=discovery_extensions)
 
     pkg_names = {d.name for d in descriptors}
-    _check_package_selection_parameters(args, pkg_names)
+    _check_package_selection_parameters(
+        args, pkg_names, selection_extensions=selection_extensions)
 
     augment_packages(
-        descriptors, additional_argument_names=additional_argument_names)
+        descriptors, additional_argument_names=additional_argument_names,
+        augmentation_extensions=augmentation_extensions)
     return descriptors
 
 
-def _check_package_selection_parameters(args, pkg_names):
-    package_selection_extensions = get_package_selection_extensions()
-    for extension in package_selection_extensions.values():
+def _check_package_selection_parameters(
+    args, pkg_names, *, selection_extensions=None,
+):
+    if selection_extensions is None:
+        selection_extensions = get_package_selection_extensions()
+    for extension in selection_extensions.values():
         try:
             retval = extension.check_parameters(args=args, pkg_names=pkg_names)
             assert retval is None, 'check_parameters() should return None'
@@ -211,19 +264,25 @@ def _check_package_selection_parameters(args, pkg_names):
             # skip failing extension, continue with next one
 
 
-def select_package_decorators(args, decorators):
+def select_package_decorators(
+    args, decorators, *, selection_extensions=None,
+):
     """
     Select the package decorators based on the command line arguments.
 
     The `selected` attribute of each decorator is updated by this function.
 
     :param args: The parsed command line arguments
     :param list decorators: The package decorators in topological order
+    :param selection_extensions: The package selection extensions to use, if
+      `None` is passed use the extensions provided by
+      :function:`get_package_selection_extensions`
     """
     # filtering must happen after the topological ordering since otherwise
     # packages in the middle of the dependency graph might be missing
-    package_selection_extensions = get_package_selection_extensions()
-    for extension in package_selection_extensions.values():
+    if selection_extensions is None:
+        selection_extensions = get_package_selection_extensions()
+    for extension in selection_extensions.values():
         try:
             retval = extension.select_packages(
                 args=args, decorators=decorators)