Repository: explosion/thinc

Issue:
Mypy plugin crash analyzing `concatenate` called with `*args`
While working on [some changes to SpaCy's type signatures](https://github.com/explosion/spaCy/pull/8773), I noticed that running Mypy on the SpaCy code base produces the following error:

```python
Traceback (most recent call last):
  File "mypy/checkexpr.py", line 3911, in accept
  File "mypy/nodes.py", line 1558, in accept
  File "mypy/checkexpr.py", line 271, in visit_call_expr
  File "mypy/checkexpr.py", line 353, in visit_call_expr_inner
  File "mypy/checkexpr.py", line 858, in check_call_expr_with_callee_type
  File "mypy/checkexpr.py", line 917, in check_call
  File "mypy/checkexpr.py", line 1029, in check_callable_call
  File "mypy/checkexpr.py", line 729, in apply_function_plugin
  File ".../spaCy/.env/lib/python3.9/site-packages/thinc/mypy.py", line 31, in function_hook
    return get_reducers_type(ctx)
  File ".../spaCy/.env/lib/python3.9/site-packages/thinc/mypy.py", line 81, in get_reducers_type
    ctx.default_return_type.type, [arg_types[0].args[0], arg_types[0].args[1]]
IndexError: tuple index out of range
```

Based on some pdb debugging (using `mypy --pdb`) and reading the Thinc Mypy plugin code, I was able to come up with this minimal test case:

```
from thinc.api import chain, concatenate, Relu, Softmax

n_hidden = 32
dropout = 0.2

model1 = chain(
    Relu(nO=n_hidden, dropout=dropout),
    Relu(nO=n_hidden, dropout=dropout),
    Softmax()
)

model2 = chain(
    Relu(nO=n_hidden, dropout=dropout),
    Relu(nO=n_hidden, dropout=dropout),
    Softmax()
)

model3 = concatenate(*[model1, model2])
```

In this example, the use of variadic arguments (`*args`) is necessary for the error to occur.

I researched Mypy's typing system and looked into the Mypy plugin code, but wasn't able to figure out what the expected behavior would be here. The root problem is that [these lines](https://github.com/explosion/thinc/blob/a125b8cf6a518bf61a418dbd077d8849e8c30167/thinc/mypy.py#L53-L56), which build the `arg_types` list, result in only a single entry in `arg_types`:

```python
[builtins.list[thinc.model.Model*[thinc.types.Floats2d, thinc.types.Floats2d]]]
```

This is because the `get_reducers_type` function doesn't take into consideration the value of `ctx.arg_kinds`:

```python
(Pdb) p ctx
FunctionContext(..., arg_kinds=[[2]], ...)
```

(based on [this lookup table](https://github.com/python/mypy/blob/9f409c3d89c7fa65eff8cf67883473487e6effd4/mypy/nodes.py#L1525-L1537), 2 means `ArgKind.ARG_STAR`)

At this point I've probably researched the issue enough to fix it myself :joy: My proposal for how to fix this is to expand variadic arguments into two "fake arguments" during type-checking (in `get_reducers_type`). This would be an approximation of all of the possible inputs for the function, so it may not be totally sound, but might be better than the current crash?

Alternatively, this could be detected as a special-case and `Model[Any]` could be returned (or some generic indicator like that) 🙂

Thanks for making thinc! 😄


Discussion:
Oh wow, nice digging! 🤩 

Your solution sounds reasonable to me, but I definitely don't know all the `mypy` ins and outs. Would love to hear what others think. 

If it's not too much additional trouble, would you be able to create the `thinc` PR? (just refer to this issue with the full explanation, ofcourse)

Resolving patch:
diff --git a/thinc/mypy.py b/thinc/mypy.py
--- a/thinc/mypy.py
+++ b/thinc/mypy.py
@@ -35,28 +35,128 @@ def function_hook(ctx: FunctionContext) -> Type:
 
 
 def get_reducers_type(ctx: FunctionContext) -> Type:
-    assert isinstance(ctx.context, CallExpr)
+    """
+    Determine a more specific model type for functions that combine models.
+
+    This function operates on function *calls*. It analyzes each function call
+    by looking at the function definition and the arguments passed as part of
+    the function call, then determines a more specific return type for the
+    function call.
+
+    This method accepts a `FunctionContext` as part of the Mypy plugin
+    interface. This function context provides easy access to:
+    * `args`: List of "actual arguments" filling each "formal argument" of the
+      called function. "Actual arguments" are those passed to the function
+      as part of the function call. "Formal arguments" are the parameters
+      defined by the function definition. The same actual argument may serve
+      to fill multiple formal arguments. In some cases the relationship may
+      even be ambiguous. For example, calling `range(*args)`, the actual
+      argument `*args` may fill the `start`, `stop` or `step` formal
+      arguments, depending on the length of the list.
+
+      The `args` list is of length `num_formals`, with each element
+      corresponding to a formal argument. Each value in the `args` list is a
+      list of actual arguments which may fill the formal argument. For
+      example, in the function call `range(*args, num)`, `num` may fill the
+      `start`, `end` or `step` formal arguments depending on the length of
+      `args`, so type-checking needs to consider all of these possibilities.
+    * `arg_types`: Type annotation (or inferred type) of each argument. Like
+      `args`, this value is a list of lists with an outer list entry for each
+      formal argument and an inner list entry for each possible actual
+      argument for the formal argument.
+    * `arg_kinds`: "Kind" of argument passed to the function call. Argument
+      kinds include positional, star (`*args`), named (`x=y`) and star2
+      (`**kwargs`) arguments (among others). Like `args`, this value is a list
+      of lists.
+    * `context`: AST node representing the function call with all available
+      type information. Notable attributes include:
+      * `args` and `arg_kinds`: Simple list of actual arguments, not mapped to
+        formal arguments.
+      * `callee`: AST node representing the function being called. Typically
+        this is a `NameExpr`. To resolve this node to the function definition
+        it references, accessing `callee.node` will usually return either a
+        `FuncDef` or `Decorator` node.
+    * etc.
+
+    This function infers a more specific type for model-combining functions by
+    making certain assumptions about how the function operates based on the
+    order of its formal arguments and its return type.
+
+    If the return type is `Model[InT, XY_YZ_OutT]`, the output of each
+    argument is expected to be used as the input to the next argument. It's
+    therefore necessary to check that the output type of each model is
+    compatible with the input type of the following model. The combined model
+    has the type `Model[InT, OutT]`, where `InT` is the input type of the
+    first model and `OutT` is the output type of the last model.
+
+    If the return type is `Model[InT, XY_XY_OutT]`, all model arguments
+    receive input of the same type and are expected to produce output of the
+    same type. It's therefore necessary to check that all models have the same
+    input types and the same output types. The combined model has the type
+    `Model[InT, OutT]`, where `InT` is the input type of all model arguments
+    and `OutT` is the output type of all model arguments.
+
+    Raises:
+        AssertionError: Raised if a more specific model type couldn't be
+            determined, indicating that the default general return type should
+            be used.
+    """
+    # Verify that we have a type-checking API and a default return type (presumably a
+    # `thinc.model.Model` instance)
     assert isinstance(ctx.api, TypeChecker)
     assert isinstance(ctx.default_return_type, Instance)
-    assert isinstance(ctx.context.callee, NameExpr)
-    assert isinstance(ctx.context.callee.node, (FuncDef, Decorator))
-    assert isinstance(ctx.context.callee.node.type, CallableType)
-    assert isinstance(ctx.context.callee.node.type.ret_type, Instance)
-    assert ctx.context.callee.node.type.ret_type.args
-    assert len(ctx.context.callee.node.type.ret_type.args) == 2
-    out_type = ctx.context.callee.node.type.ret_type.args[1]
+
+    # Verify that we're inspecting a function call to a callable defined or decorated function
+    assert isinstance(ctx.context, CallExpr)
+    callee = ctx.context.callee
+    assert isinstance(callee, NameExpr)
+    callee_node = callee.node
+    assert isinstance(callee_node, (FuncDef, Decorator))
+    callee_node_type = callee_node.type
+    assert isinstance(callee_node_type, CallableType)
+
+    # Verify that the callable returns a `thinc.model.Model`
+    # TODO: Use `map_instance_to_supertype` to map subtypes to `Model` instances.
+    # (figure out how to look up the `TypeInfo` for a class outside of the module being type-checked)
+    callee_return_type = callee_node_type.ret_type
+    assert isinstance(callee_return_type, Instance)
+    assert callee_return_type.type.fullname == thinc_model_fullname
+    assert callee_return_type.args
+    assert len(callee_return_type.args) == 2
+
+    # Obtain the output type parameter of the `thinc.model.Model` return type
+    # of the called API function
+    out_type = callee_return_type.args[1]
+
+    # Check if the `Model`'s output type parameter is one of the "special
+    # type variables" defined to represent model composition (chaining) and
+    # homogenous reduction
     assert isinstance(out_type, TypeVarType)
     assert out_type.fullname
     if out_type.fullname not in {intoin_outtoout_out_fullname, chained_out_fullname}:
         return ctx.default_return_type
+
+    # Extract type of each argument used to call the API function, making sure that they are also
+    # `thinc.model.Model` instances
     args = list(itertools.chain(*ctx.args))
     arg_types = []
     for arg_type in itertools.chain(*ctx.arg_types):
+        # TODO: Use `map_instance_to_supertype` to map subtypes to `Model` instances.
         assert isinstance(arg_type, Instance)
+        assert arg_type.type.fullname == thinc_model_fullname
+        assert len(arg_type.args) == 2
         arg_types.append(arg_type)
+
+    # Collect neighboring pairs of arguments and their types
     arg_pairs = list(zip(args[:-1], args[1:]))
     arg_types_pairs = list(zip(arg_types[:-1], arg_types[1:]))
+
+    # Determine if passed models will be chained or if they all need to have
+    # the same input and output type
     if out_type.fullname == chained_out_fullname:
+        # Models will be chained, meaning that the output of each model will
+        # be passed as the input to the next model
+        # Verify that model inputs and outputs are compatible
         for (arg1, arg2), (type1, type2) in zip(arg_pairs, arg_types_pairs):
             assert isinstance(type1, Instance)
             assert isinstance(type2, Instance)
@@ -65,10 +165,14 @@ def get_reducers_type(ctx: FunctionContext) -> Type:
             check_chained(
                 l1_arg=arg1, l1_type=type1, l2_arg=arg2, l2_type=type2, api=ctx.api
             )
+
+        # Generated model takes the first model's input and returns the last model's output
         return Instance(
             ctx.default_return_type.type, [arg_types[0].args[0], arg_types[-1].args[1]]
         )
     elif out_type.fullname == intoin_outtoout_out_fullname:
+        # Models must have the same input and output types
+        # Verify that model inputs and outputs are compatible
         for (arg1, arg2), (type1, type2) in zip(arg_pairs, arg_types_pairs):
             assert isinstance(type1, Instance)
             assert isinstance(type2, Instance)
@@ -77,9 +181,13 @@ def get_reducers_type(ctx: FunctionContext) -> Type:
             check_intoin_outtoout(
                 l1_arg=arg1, l1_type=type1, l2_arg=arg2, l2_type=type2, api=ctx.api
             )
+
+        # Generated model accepts and returns the same types as all passed models
         return Instance(
             ctx.default_return_type.type, [arg_types[0].args[0], arg_types[0].args[1]]
         )
+
+    # Make sure the default return type is returned if no branch was selected
     assert False, "Thinc mypy plugin error: it should return before this point"
 
 


Test patch:
diff --git a/thinc/tests/regression/issue519/__init__.py b/thinc/tests/regression/issue519/__init__.py
new file mode 100644
diff --git a/thinc/tests/regression/issue519/program.py b/thinc/tests/regression/issue519/program.py
new file mode 100644
--- /dev/null
+++ b/thinc/tests/regression/issue519/program.py
@@ -0,0 +1,20 @@
+from thinc.api import chain, concatenate, Relu, Softmax
+from thinc.model import Model
+from thinc.types import Floats2d
+
+n_hidden = 32
+dropout = 0.2
+
+model1 = chain(
+    Relu(nO=n_hidden, dropout=dropout),
+    Relu(nO=n_hidden, dropout=dropout),
+    Softmax()
+)
+
+model2 = chain(
+    Relu(nO=n_hidden, dropout=dropout),
+    Relu(nO=n_hidden, dropout=dropout),
+    Softmax()
+)
+
+model3: Model[Floats2d, Floats2d] = concatenate(*[model1, model2])
diff --git a/thinc/tests/regression/issue519/test_issue519.py b/thinc/tests/regression/issue519/test_issue519.py
new file mode 100644
--- /dev/null
+++ b/thinc/tests/regression/issue519/test_issue519.py
@@ -0,0 +1,25 @@
+import subprocess
+try:
+    import importlib.resources as importlib_resources
+except ImportError:
+    import importlib_resources  # type: ignore
+
+
+import pytest
+
+
+@pytest.mark.slow
+def test_issue519():
+    """
+    Test ability of Thinc mypy plugin to handle variadic arguments.
+
+    This test can take up to 45 seconds, and is thus marked as slow.
+    """
+    # Determine the name of the parent module (which contains the test program)
+    parent_module_name = __name__[:__name__.rfind(".")]
+
+    # Load test program that calls a Thinc API with variadic arguments
+    program_text = importlib_resources.read_text(parent_module_name, "program.py")
+
+    # Ask Mypy to type-check the loaded program text
+    subprocess.run(["mypy", "--command", program_text], check=True)