Merge pull request #54 from gardenlinux/common-arguments

Factor out common arguments
gardenlinux · Dec 18, 2023 · 4f2d1e2 · 4f2d1e2
2 parents dd7f530 + a9b8914
commit 4f2d1e2
Show file tree

Hide file tree

Showing 9 changed files with 40 additions and 85 deletions.
diff --git a/src/glvd/cli/client/__init__.py b/src/glvd/cli/client/__init__.py
@@ -6,3 +6,14 @@
 
 
 cli = CliRegistry()
+
+cli.add_argument(
+    '--server',
+    default='http://localhost:5000',
+    help='the server to use',
+)
+cli.add_argument(
+    '--debug',
+    action='store_true',
+    help='enable debug output',
+)
diff --git a/src/glvd/cli/client/cve.py b/src/glvd/cli/client/cve.py
@@ -21,21 +21,11 @@ class ClientCve:
     @cli.register(
         'cve',
         arguments=[
-            cli.add_argument(
+            cli.prepare_argument(
                 'cve',
                 help='the CVE to look up',
                 metavar='CVE',
             ),
-            cli.add_argument(
-                '--server',
-                default='http://localhost:5000',
-                help='the server to use',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
         ]
     )
     def run(cve: str, server: str, debug: bool) -> None:

diff --git a/src/glvd/cli/data/__init__.py b/src/glvd/cli/data/__init__.py
@@ -6,3 +6,14 @@
 
 
 cli = CliRegistry()
+
+cli.add_argument(
+    '--database',
+    default='postgresql+asyncpg:///',
+    help='the database to use, must use asyncio compatible SQLAlchemy driver',
+)
+cli.add_argument(
+    '--debug',
+    action='store_true',
+    help='enable debug output',
+)
diff --git a/src/glvd/cli/data/combine_all.py b/src/glvd/cli/data/combine_all.py
@@ -25,21 +25,7 @@
 
 class CombineAll:
     @staticmethod
-    @cli.register(
-        'combine-all',
-        arguments=[
-            cli.add_argument(
-                '--database',
-                default='postgresql+asyncpg:///',
-                help='the database to use, must use asyncio compatible SQLAlchemy driver',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
-        ]
-    )
+    @cli.register('combine-all')
     def run(database: str, debug: bool) -> None:
         logging.basicConfig(level=debug and logging.DEBUG or logging.INFO)
         engine = create_async_engine(database, echo=debug)

diff --git a/src/glvd/cli/data/combine_deb.py b/src/glvd/cli/data/combine_deb.py
@@ -32,21 +32,7 @@
 
 class CombineDeb:
     @staticmethod
-    @cli.register(
-        'combine-deb',
-        arguments=[
-            cli.add_argument(
-                '--database',
-                default='postgresql+asyncpg:///',
-                help='the database to use, must use asyncio compatible SQLAlchemy driver',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
-        ]
-    )
+    @cli.register('combine-deb')
     def run(database: str, debug: bool) -> None:
         logging.basicConfig(level=debug and logging.DEBUG or logging.INFO)
         engine = create_async_engine(database, echo=debug)

diff --git a/src/glvd/cli/data/ingest_debsec.py b/src/glvd/cli/data/ingest_debsec.py
@@ -28,28 +28,18 @@ class IngestDebsec:
     @cli.register(
         'ingest-debsec',
         arguments=[
-            cli.add_argument(
+            cli.prepare_argument(
                 'cpe_product',
                 choices=sorted(DistCpeMapper.keys()),
                 help=f'CPE product used for data, supported: {" ".join(sorted(DistCpeMapper.keys()))}',
                 metavar='CPE_PRODUCT',
             ),
-            cli.add_argument(
+            cli.prepare_argument(
                 'dir',
                 help='data directory out of https://salsa.debian.org/security-tracker-team/security-tracker',
                 metavar='DEBSEC',
                 type=Path,
             ),
-            cli.add_argument(
-                '--database',
-                default='postgresql+asyncpg:///',
-                help='the database to use, must use asyncio compatible SQLAlchemy driver',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
         ]
     )
     def run(cpe_product: str, dir: Path, database: str, debug: bool) -> None:

diff --git a/src/glvd/cli/data/ingest_debsrc.py b/src/glvd/cli/data/ingest_debsrc.py
@@ -28,33 +28,23 @@ class IngestDebsrc:
     @cli.register(
         'ingest-debsrc',
         arguments=[
-            cli.add_argument(
+            cli.prepare_argument(
                 'cpe_product',
                 choices=sorted(DistCpeMapper.keys()),
                 help=f'CPE product used for data, supported: {" ".join(sorted(DistCpeMapper.keys()))}',
                 metavar='CPE_PRODUCT',
             ),
-            cli.add_argument(
+            cli.prepare_argument(
                 'deb_codename',
                 help='codename of APT archive',
                 metavar='CODENAME',
             ),
-            cli.add_argument(
+            cli.prepare_argument(
                 'file',
                 help='uncompressed Sources file',
                 metavar='SOURCES',
                 type=Path,
             ),
-            cli.add_argument(
-                '--database',
-                default='postgresql+asyncpg:///',
-                help='the database to use, must use asyncio compatible SQLAlchemy driver',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
         ]
     )
     def run(cpe_product: str, deb_codename: str, file: Path, database: str, debug: bool) -> None:

diff --git a/src/glvd/cli/data/ingest_nvd.py b/src/glvd/cli/data/ingest_nvd.py
@@ -27,21 +27,7 @@ class IngestNvd:
     wait: int
 
     @staticmethod
-    @cli.register(
-        'ingest-nvd',
-        arguments=[
-            cli.add_argument(
-                '--database',
-                default='postgresql+asyncpg:///',
-                help='the database to use, must use asyncio compatible SQLAlchemy driver',
-            ),
-            cli.add_argument(
-                '--debug',
-                action='store_true',
-                help='enable debug output',
-            ),
-        ]
-    )
+    @cli.register('ingest-nvd')
     def run(database: str, debug: bool) -> None:
         logging.basicConfig(level=debug and logging.DEBUG or logging.INFO)
         engine = create_async_engine(database, echo=debug)

diff --git a/src/glvd/cli/registry.py b/src/glvd/cli/registry.py
@@ -6,7 +6,6 @@
 import dataclasses
 from collections.abc import (
     Callable,
-    Iterable,
 )
 
 
@@ -20,6 +19,8 @@ class CliRegistry:
     parser: argparse.ArgumentParser
     subparsers: argparse._SubParsersAction
 
+    arguments: list[_ActionWrapper]
+
     def __init__(self) -> None:
         self.parser = argparse.ArgumentParser(
             allow_abbrev=False,
@@ -29,14 +30,18 @@ def __init__(self) -> None:
         self.subparsers = self.parser.add_subparsers(
             help='sub-command help',
         )
+        self.arguments = []
 
-    def add_argument(self, *args, **kw) -> _ActionWrapper:
+    def prepare_argument(self, *args, **kw) -> _ActionWrapper:
         return _ActionWrapper(args, kw)
 
+    def add_argument(self, *args, **kw) -> None:
+        self.arguments.append(self.prepare_argument(*args, **kw))
+
     def register(
         self,
         name: str,
-        arguments: Iterable[_ActionWrapper],
+        arguments: list[_ActionWrapper] = [],
         usage: str = '%(prog)s',
         epilog: str | None = None,
     ) -> Callable:
@@ -56,7 +61,7 @@ def register(
         )
 
         for p in (parser_main, parser_sub):
-            for w in arguments:
+            for w in arguments + self.arguments:
                 p.add_argument(*w.args, **w.kw)
 
         def wrap(func: Callable) -> Callable: