Mypy 插件用于将自定义类型别名替换为 NotRequired

Question

我想编写一个 mypy 插件，以便为

NotRequired[Optional[T]]

引入类型别名。 （正如我在这个问题中发现的，不可能用普通的Python编写这个类型别名，因为在

NotRequired

定义之外不允许使用

TypedDict

。）

我的想法是定义一个通用的

Possibly

类型，如下所示：

# possibly.__init__.py

from typing import Generic, TypeVar

T = TypeVar("T")

class Possibly(Generic[T]):
    pass

然后我希望我的插件将任何出现的

Possibly[X]

替换为

NotRequired[Optional[X]]

。我尝试了以下方法：

# 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

在断点处，我知道我必须基于

mypy.types.Type

构造

arguments

的子类的实例。但我没有找到构造

NotRequired

的方法。

mypy.types

中没有对应的类型。我认为这可能是因为

typing.NotRequired

不是一个类，而是一个

typing._SpecialForm

。 （我猜这是因为

NotRequired

不影响值类型，而是影响它出现的

.__optional_keys__

的

TypedDict

的定义。）

所以，然后我想到了一个不同的策略：我可以检查

TypedDict

，看看哪些字段被标记为

Possibly

，并设置

.__optional_keys__

实例的

TypedDict

使该字段不是必需的，然后替换通过

Possibly

键入

mypy.types.UnionType(*arguments, None)

。但我没有找到

mypy.plugin.Plugin

上使用哪种方法才能将

TypedDict

放入上下文中。

所以，我被困住了。这是我第一次深入了解

mypy

的内部。你能给我一些指导如何实现我想做的事吗？

Answer 1

在 mypy 插件实现之外，首先要注意的是，您可以声明一个充当身份泛型的类型变量：

import typing as t

_T = t.TypeVar("_T")
IdentityGeneric: t.TypeAlias = _T

num: IdentityGeneric[int] = ""  # mypy: Incompatible types in assignment (expression has type "str", variable has type "int")

这意味着

Possibly

只需要

Possibly: t.TypeAlias = _T

---

您的第一次尝试（构造

mypy.types.Type

的子类实例）是正确的 - mypy 只是将其称为

mypy.types.RequiredType

，并且

NotRequired

通过构造函数指定为实例状态，如下所示：

mypy.types.RequiredType(<type>, required=False)

。

这是实现

_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)

    # 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(ctx.type.args[0])
    # Make a new instance of a union type context to represent `<type> | None`.
    union_type = mypy.types.UnionType((ctx.type.args[0], unionee_nonetype))
    union_type.set_line(ctx.type)
    # Make the `NotRequired` type context
    not_required_type = mypy.types.RequiredType(union_type, required=False)
    not_required_type.set_line(ctx.type)
    # Make mypy analyse the newly-constructed `typing.NotRequired[<type> | None]` type,
    # then return the analysed type.
    return ctx.api.analyze_type(not_required_type)

您的合规性测试的实际应用：

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 = {}  # mypy: Missing key "possibly_string" for TypedDict "Struct"  [typeddict-item]
>>> compliant_none: Struct = {"possibly_string": None}  # OK
>>> compliant_present: Struct = {"possibly_string": "a string, indeed"}  # OK

这不是一个完整的实现 - 例如，它目前允许在

Possibly

之外使用

TypedDict

，您最好避免这种情况。我在这里没有针对这种情况进行实现，因为检测到这一点需要使用 mypy 的非公开 API（即未公开的方法

mypy.plugin.TypeAnalyzerPluginInterface

），并且任何非公开的内容将来都可能会发生变化。

---

最后注意事项：

• mypy 的插件系统很强大，但也不是完全没有文档记录。为了编写插件，了解 mypy 内部结构的最简单方法是使用现有的类型构造不正确，查看错误消息中使用的字符串或字符串模式，然后尝试使用以下方法查找 mypy 的实现：字符串/模式。例如，以下是

  typing.NotRequired

  的错误用法：

  from typing_extensions import TypedDict, NotRequired
  
  class A(TypedDict):
      a: NotRequired[int, str]  # mypy: NotRequired[] must have exactly one type argument
  

  您可以在here找到此消息，这表明尽管

  typing.NotRequired

  不是一个类，但 mypy 将其建模为像任何其他泛型一样的类型，可能是因为易于分析 AST。

• 您的插件代码的组织目前是这样的：

  possibly/
    __init__.py
    plugin.py
  

  当 mypy 加载您的插件时，

  possibly.__init__

  中的任何运行时代码都会随插件一起加载，因为 mypy 在尝试加载入口点

  possibly

  时将导入

  possibly.plugin.plugin

  。所有运行时代码，包括任何可能从第三方包中提取的代码，都会在每次运行

  mypy

  时加载。我认为这是不可取的，除非你能保证你的包是轻量级的并且没有依赖项。

  事实上，当我写这篇文章时，我意识到 numpy 的 mypy 插件 (

  numpy.typing.mypy_plugin

  ) 会因为这个组织而加载 numpy（一个大库！）。

  有很多方法可以解决这个问题，而不必将插件目录与包分开 - 你必须在

  __init__

  中实现一些东西，如果 mypy 调用它，它会尝试不加载任何运行时子包。