Add kwarg support for CLIArg and subclassing config models with CLIArg fields

conflator/conflator.py

@@ -17,13 +17,17 @@
 
 
 class CLIArg:
-    def __init__(self, *args):
+    def __init__(self, *args, **kwargs):
         self.args = args
+        self.kwargs = kwargs
         self.description = None
         self.argparse_key = None
 
     def __repr__(self):
-        return f"CLIArg(args = {self.args}, description = {self.description}, argparse_key = {self.argparse_key})"
+        return (
+            f"CLIArg(args = {self.args}, kwargs = {self.kwargs}, "
+            + f"description = {self.description}, argparse_key = {self.argparse_key})"
+        )
 
 
 class EnvVar:
@@ -116,8 +120,8 @@ def _find_models(t: Type[BaseModel], seen=None) -> set[Type[BaseModel]]:
         else:
             if issubclass(t, ConfigModel) and t not in seen:
                 seen.add(t)
-                for k, v in t.__annotations__.items():
-                    Conflator._find_models(v, seen)
+                for v in t.model_fields.values():
+                    Conflator._find_models(v.annotation, seen)
 
         return seen
 
@@ -137,7 +141,7 @@ def _get_cli_args(model: Type[BaseModel], args: set[CLIArg] = None):
 
     def load(self) -> BaseModel:
         referenced_models = Conflator._find_models(self.model)
-        # rprint(referenced_config_models)
+        # rprint(referenced_models)
 
         if self.cli:
             # Find all CLI args and then parse them
@@ -157,7 +161,7 @@ def load(self) -> BaseModel:
                 )
 
             for ca in cli_args:
-                action = self.parser.add_argument(*ca.args, help=ca.description)
+                action = self.parser.add_argument(*ca.args, **ca.kwargs, help=ca.description)
                 ca.argparse_key = action.dest
 
             self.parser.add_argument(

tests/test_cli_args.py

@@ -1,13 +1,23 @@
 from unittest.mock import patch
 
 from annotated_types import Annotated
-from pydantic import Field
+from pydantic import ConfigDict, Field
 
 from conflator import CLIArg, ConfigModel, Conflator, EnvVar
 
 
 class NestedConfig(ConfigModel):
-    nested_field: str = "default_value"
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    nested_field: Annotated[str, CLIArg("--nested")] = "default_value"
+
+
+class ConfigWithCLI(ConfigModel):
+    arg1: NestedConfig = NestedConfig()
+
+
+class InheritedConfig(ConfigWithCLI):
+    pass
 
 
 class Config(ConfigModel):
@@ -22,3 +32,11 @@ def test_cli_argument_override():
         conflator = Conflator("polytope", Config, nested={})
         config = conflator.load()
         assert config.test_email == "[email protected]"
+
+
+def test_inherited_cli_arg():
+    conflator = Conflator("test", InheritedConfig)
+    cli_args = set()
+    for m in Conflator._find_models(conflator.model):
+        cli_args |= Conflator._get_cli_args(m, cli_args)
+    assert len(cli_args) == 1