Mypy Plugin for Replacing Custom TypeAlias with NotRequired

I want to write a mypy plugin in order to introduce a type alias for NotRequired[Optional[T]]. (As I found out in this question, it is not possible to write this type alias in plain python, because NotRequired is not allowed outside of a TypedDict definition.)

My idea is to define a generic Possibly type, like so:

# possibly.__init__.py

from typing import Generic, TypeVar

T = TypeVar("T")

class Possibly(Generic[T]):
    pass

I then want my plugin to replace any occurrence of Possibly[X] with NotRequired[Optional[X]]. I tried the following approach:

# possibly.plugin

from mypy.plugin import Plugin


class PossiblyPlugin(Plugin):
    def get_type_analyze_hook(self, fullname: str):
        if fullname != "possibly.Possibly":
            return
        return self._replace_possibly

    def _replace_possibly(self, ctx):
        arguments = ctx.type.args
        breakpoint()


def plugin(version):
    return PossiblyPlugin

At the breakpoint, I understand I have to construct an instance of a subclass of mypy.types.Type based on arguments. But I didn't find a way to construct NotRequired. There is no corresponding type in mypy.types. I figure this might be due to the fact that typing.NotRequired is not a class, but a typing._SpecialForm. (I guess this is because NotRequired does not affect the value type, but the definition of the .__optional_keys__ of the TypedDict it occurs on.)

So, then I thought about a different strategy: I could check for TypedDicts, see which fields are marked Possibly, and set the .__optional_keys__ of the TypedDict instance to make the field not required, and replace the Possibly type by mypy.types.UnionType(*arguments, None). But I didn't find which method on mypy.plugin.Plugin to use in order to get the TypedDicts into the context.

So, I am stuck. It is the first time I dig into the internals of mypy. Could you give me some direction how to achieve what I want to do?

Solution

Your first attempt (to construct an instance of subclass of mypy.types.Type) was correct - mypy simply calls it mypy.types.RequiredType, and NotRequired is specified as an instance state through the constructor like this: mypy.types.RequiredType(<type>, required=False).

Here's an initial attempt at an implementation of _replace_possibly:

def _replace_possibly(ctx: mypy.plugin.AnalyzeTypeContext) -> mypy.types.Type:
    """
    Transform `possibly.Possibly[<type>]` into `typing.NotRequired[<type> | None]`. Most
    of the implementation is copied from
    `mypy.typeanal.TypeAnalyser.try_analyze_special_unbound_type`.

    All `set_line` calls in the implementation are for reporting purposes, so that if
    any errors occur, mypy will report them in the correct line and column in the file.
    """

    if len(ctx.type.args) != 1:
        ctx.api.fail(
            "possibly.Possibly[] must have exactly one type argument",
            ctx.type,
            code=mypy.errorcodes.VALID_TYPE,
        )
        return mypy.types.AnyType(mypy.types.TypeOfAny.from_error)

    # Disallow usage of `Possibly` outside of `TypedDict`. Note: This check uses
    # non-exposed API, but must be done, because (as of mypy==1.8.0) the plugin will
    # otherwise crash.
    type_analyser: mypy.typeanal.TypeAnalyser = ctx.api  # type: ignore[assignment]
    if not type_analyser.allow_required:
        ctx.api.fail(
            "possibly.Possibly[] can only be used in a TypedDict definition",
            ctx.type,
            code=mypy.errorcodes.VALID_TYPE,
        )
        return mypy.types.AnyType(mypy.types.TypeOfAny.from_error)

    # Make mypy analyse `<type>` and get the analysed type
    analysed_type = ctx.api.analyze_type(ctx.type.args[0])
    # Make a new instance of a `None` type context to represent `None` in the union
    # `<type> | None`
    unionee_nonetype = mypy.types.NoneType()
    unionee_nonetype.set_line(analysed_type)
    # Make a new instance of a union type context to represent `<type> | None`.
    union_type = mypy.types.UnionType((analysed_type, unionee_nonetype))
    union_type.set_line(ctx.type)
    # Make the `NotRequired[<type> | None]` type context
    not_required_type = mypy.types.RequiredType(union_type, required=False)
    not_required_type.set_line(ctx.type)
    return not_required_type

Your compliance tests in action:

import typing_extensions as t

import possibly

class Struct(t.TypedDict):
    possibly_string: possibly.Possibly[str]

>>> non_compliant: Struct = {"possibly_string": int}  # mypy: Incompatible types (expression has type "type[int]", TypedDict item "possibly_string" has type "str | None") [typeddict-item]
>>> compliant_absent: Struct = {}  # OK
>>> compliant_none: Struct = {"possibly_string": None}  # OK
>>> compliant_present: Struct = {"possibly_string": "a string, indeed"}  # OK

Notes:

mypy's plugin system is powerful but not thoroughly undocumented. The easiest way to go through mypy's internals, for the purposes of writing a plugin, is to use an existing type construct incorrectly, look at what string or string pattern is used in the error message, then attempt to find mypy's implementation using the string/pattern. For example, the following is an incorrect usage of typing.NotRequired:
```
from typing_extensions import TypedDict, NotRequired

class A(TypedDict):
    a: NotRequired[int, str]  # mypy: NotRequired[] must have exactly one type argument
```
You can find this message here, which indicates that, despite typing.NotRequired not being a class, mypy models it as a type like any other generic, possibly because of ease of analysing the AST.
Your plugin code's organisation is currently this:
```
possibly/
  __init__.py
  plugin.py
```
When mypy loads your plugin, any runtime code in possibly.__init__ is loaded with the plugin, because mypy will import possibly when it tries to load the entry point possibly.plugin.plugin. All runtime code, including any which may be pulled from third-party packages, will be loaded every time mypy is run. I don't think this is desirable unless you can guarantee that your package is lightweight and has no dependencies.

In fact, as I'm writing this, I realised that numpy's mypy plugin (numpy.typing.mypy_plugin) loads numpy (a big library!) because of this organisation.

There are ways around this without having to separate the plugin directory from the package - you'll have to implement something in __init__ which tries to not load any runtime subpackages if it's called by mypy.