Add support for attrs.fields

Author: Tinche

mypy/plugins/attrs.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Iterable, List, cast
+from typing import Iterable, List, Optional, cast
 from typing_extensions import Final, Literal
 
 import mypy.plugin  # To avoid circular imports.
@@ -43,7 +43,7 @@
     Var,
     is_class_var,
 )
-from mypy.plugin import SemanticAnalyzerPluginInterface
+from mypy.plugin import FunctionContext, SemanticAnalyzerPluginInterface
 from mypy.plugins.common import (
     _get_argument,
     _get_bool_argument,
@@ -988,3 +988,27 @@ def evolve_function_sig_callback(ctx: mypy.plugin.FunctionSigContext) -> Callabl
         ret_type=inst_type,
         name=f"{ctx.default_signature.name} of {inst_type_str}",
     )
+
+
+def _get_cls_from_init(t: Type) -> Optional[TypeInfo]:
+    if isinstance(t, CallableType):
+        return t.type_object()
+    return None
+
+
+def fields_function_callback(ctx: FunctionContext) -> Type:
+    """Provide the proper return value for `attrs.fields`."""
+    if ctx.arg_types and ctx.arg_types[0] and ctx.arg_types[0][0]:
+        first_arg_type = ctx.arg_types[0][0]
+        cls = _get_cls_from_init(first_arg_type)
+        if cls is not None:
+            if MAGIC_ATTR_NAME in cls.names:
+                # This is a proper attrs class.
+                ret_type = cls.names[MAGIC_ATTR_NAME].type
+                return ret_type
+            else:
+                ctx.api.fail(
+                    f'Argument 1 to "fields" has incompatible type "{format_type_bare(first_arg_type)}"; expected an attrs class',
+                    ctx.context,
+                )
+    return ctx.default_return_type

mypy/plugins/default.py

@@ -39,12 +39,14 @@ class DefaultPlugin(Plugin):
     """Type checker plugin that is enabled by default."""
 
     def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
-        from mypy.plugins import ctypes, singledispatch
+        from mypy.plugins import attrs, ctypes, singledispatch
 
         if fullname == "ctypes.Array":
             return ctypes.array_constructor_callback
         elif fullname == "functools.singledispatch":
             return singledispatch.create_singledispatch_function_callback
+        elif fullname in ("attr.fields", "attrs.fields"):
+            return attrs.fields_function_callback
         return None
 
     def get_function_signature_hook(

test-data/unit/check-attr.test

@@ -1548,6 +1548,33 @@ takes_attrs_cls(A(1, ""))  # E: Argument 1 to "takes_attrs_cls" has incompatible
 takes_attrs_instance(A)  # E: Argument 1 to "takes_attrs_instance" has incompatible type "Type[A]"; expected "AttrsInstance" # N: ClassVar protocol member AttrsInstance.__attrs_attrs__ can never be matched by a class object
 [builtins fixtures/attr.pyi]
 
+[case testAttrsFields]
+import attr
+from attrs import fields
+
+@attr.define
+class A:
+    b: int
+    c: str
+
+reveal_type(fields(A))  # N: Revealed type is "Tuple[attr.Attribute[builtins.int], attr.Attribute[builtins.str], fallback=__main__.A.____main___A_AttrsAttributes__]"
+reveal_type(fields(A)[0])  # N: Revealed type is "attr.Attribute[builtins.int]"
+reveal_type(fields(A).b)  # N: Revealed type is "attr.Attribute[builtins.int]"
+fields(A).x  # E: "____main___A_AttrsAttributes__" has no attribute "x"
+
+[builtins fixtures/attr.pyi]
+
+[case testNonattrsFields]
+from attrs import fields
+
+class A:
+    b: int
+    c: str
+
+fields(A)  # E: Argument 1 to "fields" has incompatible type "Type[A]"; expected an attrs class
+
+[builtins fixtures/attr.pyi]
+
 [case testAttrsInitMethodAlwaysGenerates]
 from typing import Tuple
 import attr

test-data/unit/lib-stub/attr/__init__.pyi

@@ -247,3 +247,5 @@ def field(
 
 def evolve(inst: _T, **changes: Any) -> _T: ...
 def assoc(inst: _T, **changes: Any) -> _T: ...
+
+def fields(cls: _C) -> Any: ...

test-data/unit/lib-stub/attrs/__init__.pyi

@@ -129,3 +129,5 @@ def field(
 
 def evolve(inst: _T, **changes: Any) -> _T: ...
 def assoc(inst: _T, **changes: Any) -> _T: ...
+
+def fields(cls: _C) -> Any: ...