From 32e29bd61e2943c08ded7a23606c53096a57c7a1 Mon Sep 17 00:00:00 2001
From: Pedro Brochado <pedropsb95@gmail.com>
Date: Tue, 20 Feb 2024 09:29:26 -0300
Subject: [PATCH] Fix support for subpackages

---
 src/pulp_docs/mkdocs_macros.py     | 24 ++++++++++++------------
 src/pulp_docs/repository.py        | 14 ++++++--------
 src/pulp_docs/utils/aggregation.py |  3 +--
 3 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/src/pulp_docs/mkdocs_macros.py b/src/pulp_docs/mkdocs_macros.py
index cf3456c..62a68cf 100644
--- a/src/pulp_docs/mkdocs_macros.py
+++ b/src/pulp_docs/mkdocs_macros.py
@@ -90,23 +90,23 @@ def prepare_repositories(TMPDIR: Path, repos: Repos, config: Config):
     shutil.rmtree(repo_sources, ignore_errors=True)
     shutil.rmtree(repo_docs, ignore_errors=True)
 
-    for repo in repos.all:
+    for repo_or_pkg in repos.all:
         start = time.perf_counter()
         # handle subpcakges nested under repositories
-        this_docs_dir = repo_docs / repo.name
-        if not isinstance(repo, SubPackage):
-            this_src_dir = repo_sources / repo.name
-            repo.download(dest_dir=this_src_dir, clear_cache=config.clear_cache)
+        this_docs_dir = repo_docs / repo_or_pkg.name
+        if not isinstance(repo_or_pkg, SubPackage):
+            this_src_dir = repo_sources / repo_or_pkg.name
+            repo_or_pkg.download(dest_dir=this_src_dir, clear_cache=config.clear_cache)
         else:
-            this_src_dir = repo_sources / repo.subpackage_of / repo.name
+            this_src_dir = repo_sources / repo_or_pkg.subpackage_of / repo_or_pkg.name
 
         # install and post-process
-        _install_doc_files(this_src_dir, this_docs_dir, repo)
-        if repo.type == "content":
-            _generate_rest_api_page(this_docs_dir, repo.name, repo.title)
+        _install_doc_files(this_src_dir, this_docs_dir, repo_or_pkg)
+        if repo_or_pkg.type == "content":
+            _generate_rest_api_page(this_docs_dir, repo_or_pkg.name, repo_or_pkg.title)
 
         end = time.perf_counter()
-        log.info(f"{repo.name} completed in {end - start:.2} sec")
+        log.info(f"{repo_or_pkg.name} completed in {end - start:.2} sec")
 
     # Copy core-files (shipped with pulp-docs) to tmpdir
     shutil.copy(
@@ -116,8 +116,8 @@ def prepare_repositories(TMPDIR: Path, repos: Repos, config: Config):
 
     # Log
     log.info("[pulp-docs] Done downloading sources. Here are the sources used:")
-    for repo in repos.all:
-        log.info({repo.name: str(repo.status)})
+    for repo_or_pkg in repos.all:
+        log.info({repo_or_pkg.name: str(repo_or_pkg.status)})
 
     return (repo_docs, repo_sources)
 
diff --git a/src/pulp_docs/repository.py b/src/pulp_docs/repository.py
index dea7642..2484201 100644
--- a/src/pulp_docs/repository.py
+++ b/src/pulp_docs/repository.py
@@ -239,17 +239,15 @@ def all(self):
                 subpackages.extend(repo.subpackages)
         return repos + subpackages
 
-    def get_repos(self, type="all"):
+    def get_repos(self, repo_types: t.Optional[t.List] = None):
         """Get a set of repositories and subpackages by type."""
-        if type == "all":
+        # Default case
+        if repo_types is None:
             return self.all
 
-        repos_and_pkgs = []
-        for repo in self.all:
-            if getattr(repo, "subpackages", None):
-                repos_and_pkgs.extend(repo.subpackages)
-
-        return [repo for repo in repos_and_pkgs if repo.type == type]
+        # Filter by repo_types
+        repos_and_pkgs = self.all
+        return [repo for repo in repos_and_pkgs if repo.type in repo_types]
 
     @classmethod
     def from_yaml(cls, path: str):
diff --git a/src/pulp_docs/utils/aggregation.py b/src/pulp_docs/utils/aggregation.py
index a052c0b..39ec2d3 100644
--- a/src/pulp_docs/utils/aggregation.py
+++ b/src/pulp_docs/utils/aggregation.py
@@ -93,8 +93,7 @@ def repo_grouping(
         if not repo_types:  # default case
             selected_repos = self.repos.all
         else:
-            for repo_name in repo_types:
-                selected_repos.extend(self.repos.get_repos(type=type))
+            selected_repos.extend(self.repos.get_repos(repo_types=repo_types))
 
         # Dont expand content-types
         if not _expand_content_types: