Allow skipping the extraction of organization members and/or repositories. (#13)

* feat: allow skipping org repos or members

* chore: refactor args

Author: jbristow

nodestream_github/orgs.py

@@ -19,8 +19,18 @@
 
 
 class GithubOrganizationsExtractor(Extractor):
-    def __init__(self, **github_client_kwargs: any):
-        self.client = GithubRestApiClient(**github_client_kwargs)
+    def __init__(
+        self,
+        *,
+        include_members: bool | None = True,
+        include_repositories: bool | None = True,
+        **kwargs: any,
+    ):
+
+        self.include_members = include_members is True
+        self.include_repositories = include_repositories is True
+
+        self.client = GithubRestApiClient(**kwargs)
 
     async def extract_records(self) -> AsyncGenerator[OrgRecord]:
         async for org in self.client.fetch_all_organizations():
@@ -34,14 +44,23 @@ async def _extract_organization(self, login: str) -> OrgRecord | None:
         if not full_org:
             return None
 
-        full_org["members"] = [user async for user in self._fetch_all_members(login)]
+        if self.include_members:
+            full_org["members"] = [
+                user async for user in self._fetch_all_members(login)
+            ]
+        else:
+            full_org["members"] = []
+
+        if self.include_repositories:
+            full_org["repositories"] = [
+                simplify_repo(
+                    repo, permission=full_org.get("default_repository_permission")
+                )
+                async for repo in self.client.fetch_repos_for_org(login)
+            ]
+        else:
+            full_org["repositories"] = []
 
-        full_org["repositories"] = [
-            simplify_repo(
-                repo, permission=full_org.get("default_repository_permission")
-            )
-            async for repo in self.client.fetch_repos_for_org(login)
-        ]
         return full_org
 
     async def _fetch_all_members(self, login: str) -> AsyncGenerator[SimplifiedUser]:

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nodestream-plugin-github"
-version = "0.13.1-beta.8"
+version = "0.13.1-beta.9"
 description = ""
 authors = ["Jon Bristow <[email protected]>"]
 packages = [

pyproject.toml

@@ -268,3 +268,77 @@ async def test_get_orgs(
             }],
         }
     ]
+
+
+@pytest.mark.asyncio
+async def test_skip_members(
+    gh_rest_mock: GithubHttpxMock,
+):
+    org_client = GithubOrganizationsExtractor(
+        auth_token="test-token",
+        github_hostname=DEFAULT_HOSTNAME,
+        user_agent="test-agent",
+        max_retries=0,
+        per_page=DEFAULT_PER_PAGE,
+        include_members=False,
+    )
+
+    gh_rest_mock.all_orgs(json=[GITHUB_ORG_SUMMARY])
+    gh_rest_mock.get_org("github", json=GITHUB_ORG)
+    gh_rest_mock.get_repos_for_org("github", json=[HELLO_WORLD_REPO])
+
+    all_records = [record async for record in org_client.extract_records()]
+    assert all_records == [
+        BASE_EXPECTED_GITHUB_ORG
+        | {
+            "members": [],
+            "repositories": [{
+                "full_name": "octocat/Hello-World",
+                "html_url": "https://github.com/octocat/Hello-World",
+                "id": 1296269,
+                "name": "Hello-World",
+                "node_id": "MDEwOlJlcG9zaXRvcnkxMjk2MjY5",
+                "permission": "read",
+                "url": "https://HOSTNAME/repos/octocat/Hello-World",
+            }],
+        }
+    ]
+
+
+@pytest.mark.asyncio
+async def test_skip_repositories(gh_rest_mock: GithubHttpxMock):
+    org_client = GithubOrganizationsExtractor(
+        auth_token="test-token",
+        github_hostname=DEFAULT_HOSTNAME,
+        include_repositories=False,  # putting the here to test kwargs interaction
+        user_agent="test-agent",
+        max_retries=0,
+        per_page=DEFAULT_PER_PAGE,
+    )
+
+    gh_rest_mock.all_orgs(json=[GITHUB_ORG_SUMMARY])
+    gh_rest_mock.get_org("github", json=GITHUB_ORG)
+    gh_rest_mock.get_members_for_org("github", json=[OCTOCAT_USER], role="admin")
+    gh_rest_mock.get_members_for_org("github", json=[TURBO_USER], role="member")
+
+    all_records = [record async for record in org_client.extract_records()]
+    assert all_records == [
+        BASE_EXPECTED_GITHUB_ORG
+        | {
+            "members": [
+                {
+                    "id": 1,
+                    "login": "octocat",
+                    "node_id": "MDQ6VXNlcjE=",
+                    "role": "admin",
+                },
+                {
+                    "id": 2,
+                    "login": "turbo",
+                    "node_id": "MDQ6VXNlcjI=",
+                    "role": "member",
+                },
+            ],
+            "repositories": [],
+        }
+    ]