Add support for custom Option and Argument subclasses (#16)

Author: uhlajs

dataclass_click/dataclass_click.py

@@ -185,13 +185,9 @@ def _patch_click_types(
         hint: typing.Type[Any]
         _, hint = _strip_optional(type_hints[key])
         if "type" not in annotation.kwargs:
-            stub: click.core.Option | click.core.Argument
-            if annotation.callable is click.option:
-                stub = click.core.Option(annotation.args, **annotation.kwargs)
-                if stub.is_flag:
-                    continue
-            else:
-                stub = click.core.Argument(annotation.args, **annotation.kwargs)
+            stub = _build_stub(annotation)
+            if hasattr(stub, "is_flag") and stub.is_flag:
+                continue
             annotation.kwargs["type"] = _eval_type(key, hint, stub, complete_type_inferences)
 
 
@@ -218,6 +214,23 @@ def _eval_type(
     raise TypeError(f"Could not infer ParamType for {key} type {hint!r}. Explicitly annotate type=<type>")
 
 
+def _build_stub(annotation: _DelayedCall) -> click.core.Option | click.core.Argument:
+    """Stub uses click's parser rather than trying to second guess how click will behave."""
+    if "cls" in annotation.kwargs:
+        return annotation.kwargs["cls"](
+            annotation.args,
+            **{k: v
+               for k, v in annotation.kwargs.items() if k != "cls"},
+        )
+
+    if annotation.callable is click.option:
+        return click.core.Option(annotation.args, **annotation.kwargs)
+    elif annotation.callable is click.argument:
+        return click.core.Argument(annotation.args, **annotation.kwargs)
+
+    raise TypeError(f"Unknown annotation callable {annotation.callable!r}")
+
+
 def _patch_required(arg_class: typing.Type[Arg], annotations: dict[str, _DelayedCall]) -> None:
     """Default click option to required if typehint is not OPTIONAL
 
@@ -228,15 +241,15 @@ def _patch_required(arg_class: typing.Type[Arg], annotations: dict[str, _Delayed
     :return: None, annotations are updated in place"""
     type_hints = typing.get_type_hints(arg_class)
     for key, annotation in annotations.items():
-        hint: typing.Type[Any]
-        is_optional, hint = _strip_optional(type_hints[key])
+        is_optional, _ = _strip_optional(type_hints[key])
         if not is_optional:
             if annotation.callable is click.option:
                 # If required or default set directly.
                 if "required" not in annotation.kwargs and "default" not in annotation.kwargs:
-                    # Stub uses click's parser rather than trying to second guess how click will behave
+                    stub = _build_stub(annotation)
+                    if not isinstance(stub, click.core.Option):
+                        raise TypeError(f"Expected click.core.Option or its subclass, got {stub!r}")
                     # If click would imply is_flag or multiple
-                    stub = click.core.Option(annotation.args, **annotation.kwargs)
                     if not stub.is_flag and not stub.multiple:
                         annotation.kwargs["required"] = True
 

tests/mypy_check.py

@@ -65,6 +65,7 @@ def main_d(a: Config, b: Config2):
 def main_e(a: Config, b: Config2):
     ...
 
+
 @click.command()
 @dataclass_click(Config2)
 @dataclass_click(Config)

tests/test_end_to_end.py

@@ -315,3 +315,57 @@ def test_mypy_check_loads():
     """Check that mypy_check.py will parse by python"""
     from . import mypy_check
     assert mypy_check
+
+
+def test_cls_kwarg():
+
+    class MutexOption(click.Option):
+
+        def __init__(self, *args, **kwargs):
+            self.not_required_if: list[str] = kwargs.pop("not_required_if")
+
+            assert self.not_required_if, "'not_required_if' parameter required"
+            kwargs["help"] = (
+                f"{kwargs.get('help', '')} Option is mutually exclusive with "
+                f"{', '.join(self.not_required_if)}.".strip())
+            super().__init__(*args, **kwargs)
+
+        def handle_parse_result(self, ctx, opts, args):
+            if set(self.not_required_if).intersection(set(opts)):
+                self.required = False
+
+            return super().handle_parse_result(ctx, opts, args)
+
+    class MutexArgument(click.Argument):
+
+        def __init__(self, *args, **kwargs):
+            self.not_required_if: list[str] = kwargs.pop("not_required_if")
+
+            assert self.not_required_if, "'not_required_if' parameter required"
+            super().__init__(*args, **kwargs)
+
+        def handle_parse_result(self, ctx, opts, args):
+            if set(self.not_required_if).intersection(set(opts)):
+                self.required = False
+
+            return super().handle_parse_result(ctx, opts, args)
+
+    @dataclass
+    class Config:
+        username: Annotated[str | None, option(required=True, cls=MutexOption, not_required_if=["token"])]
+        password: Annotated[str | None, option(required=True, cls=MutexOption, not_required_if=["token"])]
+        token: Annotated[str | None,
+                         argument(required=True, cls=MutexArgument, not_required_if=["username", "password"])]
+
+    @click.command()
+    @dataclass_click(Config)
+    def main(*args, **kwargs):
+        results.append((args, kwargs))
+
+    results: list[CallRecord] = []
+    quick_run(main, "--username", "foo", "--password", "bar")
+    assert results == [((Config(username="foo", password="bar", token=None), ), {})]
+
+    results: list[CallRecord] = []
+    quick_run(main, "foo")
+    assert results == [((Config(username=None, password=None, token="foo"), ), {})]