From 072410c52188e0e6cc071fa99027ae38125608d6 Mon Sep 17 00:00:00 2001
From: Felix Fontein <felix@fontein.de>
Date: Sat, 3 Oct 2020 12:28:54 +0200
Subject: [PATCH] Improve filtering.

---
 antsibull/data/collection-enum.py          | 41 ++++++++++++++++------
 antsibull/docs_parsing/ansible_internal.py |  6 ++--
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/antsibull/data/collection-enum.py b/antsibull/data/collection-enum.py
index 2335a7df..9e7ec266 100644
--- a/antsibull/data/collection-enum.py
+++ b/antsibull/data/collection-enum.py
@@ -77,6 +77,19 @@ def load_plugin(loader, plugin_type, plugin):
     return result
 
 
+def ansible_doc_coll_filter(coll_filter):
+    return coll_filter[0] if coll_filter and len(coll_filter) == 1 else None
+
+
+def match_filter(name, coll_filter):
+    if coll_filter is None or name in coll_filter:
+        return True
+    for filter in coll_filter:
+        if name.startswith(filter + '.'):
+            return True
+    return False
+
+
 def load_all_plugins(plugin_type, basedir, coll_filter):
     loader = getattr(plugin_loader, '%s_loader' % plugin_type)
 
@@ -87,17 +100,19 @@ def load_all_plugins(plugin_type, basedir, coll_filter):
 
     plugin_list = set()
 
-    if coll_filter is None:
+    if match_filter('ansible.builtin', coll_filter):
         paths = loader._get_paths_with_context()
         for path_context in paths:
             plugin_list.update(
                 doc.DocCLI.find_plugins(path_context.path, path_context.internal, plugin_type))
 
-    doc.add_collection_plugins(plugin_list, plugin_type, coll_filter=coll_filter)
+    doc.add_collection_plugins(
+        plugin_list, plugin_type, coll_filter=ansible_doc_coll_filter(coll_filter))
 
     result = {}
     for plugin in plugin_list:
-        result[plugin] = load_plugin(loader, plugin_type, plugin)
+        if match_filter(plugin, coll_filter):
+            result[plugin] = load_plugin(loader, plugin_type, plugin)
 
     return result
 
@@ -105,7 +120,7 @@ def load_all_plugins(plugin_type, basedir, coll_filter):
 def main(args):
     parser = argparse.ArgumentParser(
         prog=args[0], description='Bulk extraction of Ansible plugin docs.')
-    parser.add_argument('args', nargs='?', help='Collection filter', metavar='collection_filter')
+    parser.add_argument('args', nargs='*', help='Collection filter', metavar='collection_filter')
     parser.add_argument('--pretty', action='store_true', help='Pretty-print JSON')
     opt_help.add_basedir_options(parser)
 
@@ -127,17 +142,21 @@ def main(args):
         result['plugins'][plugin_type] = load_all_plugins(plugin_type, basedir, coll_filter)
 
     # Export collection data
-    b_colldirs = list_collection_dirs(coll_filter=coll_filter)
+    b_colldirs = list_collection_dirs(coll_filter=ansible_doc_coll_filter(coll_filter))
     for b_path in b_colldirs:
         collection = CollectionRequirement.from_path(b_path, False, fallback_metadata=True)
 
-        result['collections']['{0}.{1}'.format(collection.namespace, collection.name)] = {
-            'path': to_native(b_path),
-            'version': collection.metadata.version if collection.metadata.version != '*' else None,
+        collection_name = '{0}.{1}'.format(collection.namespace, collection.name)
+        if match_filter(collection_name, coll_filter):
+            version = collection.metadata.version
+            result['collections'][collection_name] = {
+                'path': to_native(b_path),
+                'version': version if version != '*' else None,
+            }
+    if match_filter('ansible.builtin', coll_filter):
+        result['collections']['ansible.builtin'] = {
+            'version': ansible_release.__version__,
         }
-    result['collections']['ansible.builtin'] = {
-        'version': ansible_release.__version__,
-    }
 
     print(json.dumps(
         result, cls=AnsibleJSONEncoder, sort_keys=True, indent=4 if arguments.pretty else None))
diff --git a/antsibull/docs_parsing/ansible_internal.py b/antsibull/docs_parsing/ansible_internal.py
index bc0b9baa..985ad7d1 100644
--- a/antsibull/docs_parsing/ansible_internal.py
+++ b/antsibull/docs_parsing/ansible_internal.py
@@ -50,9 +50,9 @@ async def get_ansible_plugin_info(venv: t.Union['VenvRunner', 'FakeVenvRunner'],
     with tempfile.NamedTemporaryFile() as tmp_file:
         tmp_file.write(get_antsibull_data('collection-enum.py'))
         collection_enum_args = [tmp_file.name]
-        if collection_names and len(collection_names) == 1:
-            # Script allows to filter by one collection
-            collection_enum_args.append(collection_names[0])
+        if collection_names is not None:
+            # Script allows to filter by collections
+            collection_enum_args.extend(collection_names)
         collection_enum_cmd = venv_python(*collection_enum_args, _env=env)
         raw_result = collection_enum_cmd.stdout.decode('utf-8', errors='surrogateescape')
         result = json.loads(_filter_non_json_lines(raw_result)[0])