Pyright Incorrect Type Inference With NewType And Generics

Hey guys! Let's dive into a quirky issue in Pyright where type inference goes a bit wonky when using NewType with generics. It's like when you expect one thing, but the code gives you something totally different. This can be a head-scratcher, so let’s break it down and see what’s happening.

The Curious Case of NewType and Generics

So, the core of the problem lies in how Pyright handles NewType when it's used with generics. Imagine you're building an application with a dependency injection (DI) container. You might have an API that looks something like this:

class Container:
    def get(self, dependency_type: type[T]) -> T: ...

Here, you're trying to fetch a dependency of a certain type T. Now, if you use NewType as the dependency_type, things get a little weird. It seems like Pyright isn't quite inferring the return type correctly. Let's see a real-world example to make this clearer.

Diving into the Code Example

Consider this scenario. You define a NewType called TestType based on an integer:

import typing as t

TestType = t.NewType("TestType", int)
T = t.TypeVar("T")

Now, let's say you have a Container class with a get method that's supposed to return an instance of the type you pass in:

class Container:
    def get(self, dependency_type: type[T]) -> T:
        return dependency_type()

If you try to use this with TestType, you might run into an issue. When you call container.get(TestType), Pyright might not infer the type as TestType. Instead, it might throw an error or infer an incorrect type, like FunctionType.

The Error Message

When this happens, you might see an error message like this from the Pyright CLI:

Argument of type "type[TestType]" cannot be assigned to parameter "dependency_type" of type "type[T@get]" in function "get"
      Type "FunctionType" is not assignable to type "type[T@get]" (reportArgumentType)

This error is essentially saying that Pyright is having trouble matching the type you passed (TestType) with the expected generic type T in the get method. It's as if Pyright is seeing TestType more as a function type than the actual NewType you defined.

Another Tricky Case

To make things even more interesting, there's another scenario where this issue pops up. Consider a simple function foo that takes a generic type T and returns it:

def foo(x: T) -> T:
    return x

If you call foo with an instance of TestType, like foo(TestType(42)), Pyright correctly infers the return type as TestType. However, if you pass TestType itself, like foo(TestType), Pyright infers the type as FunctionType. This is quite unexpected because you’d think passing the type TestType should still result in TestType being inferred.

Full Code Example

Here’s the complete code snippet that demonstrates these issues:

import typing as t


TestType = t.NewType("TestType", int)
T = t.TypeVar("T")


class Container:
    def get(self, dependency_type: type[T]) -> T:
        return dependency_type()


def foo(x: T) -> T:
    return x


container = Container()


t.reveal_type(container.get(TestType))
t.reveal_type(foo(TestType(5)))
t.reveal_type(foo(TestType))

In this example, t.reveal_type is used to show the inferred types. You’ll notice that container.get(TestType) and foo(TestType) don’t behave as you might expect.

Why Does This Happen?

So, why is Pyright acting this way? It seems to stem from how Pyright resolves generic types when NewType is involved. NewType creates a distinct type, but it's closely related to its parent type (in this case, int). When you pass TestType as a type argument, Pyright might be getting confused about whether you're passing the type itself or an instance of the type. This confusion leads to the incorrect inference of FunctionType.

The issue might be rooted in the intricacies of typeVar resolution and how Pyright handles the dual nature of NewType – it's both a type and a callable (since you can use it to create instances).

Mypy's Perspective

Interestingly, this behavior isn't consistent across all type checkers. Mypy, for instance, doesn’t throw errors in these scenarios. This suggests that the issue might be specific to Pyright's type inference algorithm and how it interacts with NewType and generics.

Possible Workarounds and Solutions

Now that we understand the problem, let's think about how we can work around it. While we wait for a potential fix in Pyright, there are a few strategies we can use.

Explicit Type Casting

One straightforward approach is to use explicit type casting. This involves telling Pyright exactly what type you expect. For example:

result = t.cast(TestType, container.get(TestType))
t.reveal_type(result)  # Revealed type is "TestType"

By using t.cast, you're essentially telling Pyright,