Extend mypy plugin (#290)

supriya-project · Dec 5, 2022 · b79e0e1 · b79e0e1
1 parent c0ecd95
commit b79e0e1
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/supriya/ext/mypy.py b/supriya/ext/mypy.py
@@ -87,12 +87,18 @@ def transform(self) -> bool:
         for name in ["ar", "kr", "ir", "dr", "new"]:
             if not decorator_arguments[name] or name in info.names:
                 continue
+            return_type = api.named_type("supriya.ugens.bases.UGenMethodMixin")
+            if (
+                decorator_arguments["is_multichannel"]
+                or decorator_arguments["fixed_channel_count"]
+            ):
+                return_type = api.named_type("supriya.ugens.bases.UGenArray")
             add_method_to_class(
                 api=api,
                 cls=cls,
                 name=name,
                 args=args,
-                return_type=api.named_type("supriya.ugens.bases.UGenMethodMixin"),
+                return_type=return_type,
                 is_classmethod=True,
             )
         if "__init__" not in info.names: