Add __attrs_attrs__ and fields to attrs

mypy/plugins/attrs.py

@@ -1,5 +1,6 @@
 """Plugin for supporting the attrs library (http://www.attrs.org)"""
 
+from mypy.lookup import lookup_fully_qualified
 from mypy.ordered_dict import OrderedDict
 
 from typing import Optional, Dict, List, cast, Tuple, Iterable
@@ -9,19 +10,19 @@
 from mypy.exprtotype import expr_to_unanalyzed_type, TypeTranslationError
 from mypy.fixup import lookup_qualified_stnode
 from mypy.nodes import (
-    Context, Argument, Var, ARG_OPT, ARG_POS, TypeInfo, AssignmentStmt,
-    TupleExpr, ListExpr, NameExpr, CallExpr, RefExpr, FuncDef,
+    Block, ClassDef, Context, Argument, SymbolTable, Var, ARG_OPT, ARG_POS, TypeInfo,
+    AssignmentStmt, TupleExpr, ListExpr, NameExpr, CallExpr, RefExpr, FuncDef,
     is_class_var, TempNode, Decorator, MemberExpr, Expression,
     SymbolTableNode, MDEF, JsonDict, OverloadedFuncDef, ARG_NAMED_OPT, ARG_NAMED,
     TypeVarExpr, PlaceholderNode
 )
-from mypy.plugin import SemanticAnalyzerPluginInterface
+from mypy.plugin import FunctionContext, SemanticAnalyzerPluginInterface
 from mypy.plugins.common import (
     _get_argument, _get_bool_argument, _get_decorator_bool_argument, add_method,
     deserialize_and_fixup_type
 )
 from mypy.types import (
-    Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarDef, TypeVarType,
+    Instance, Type, AnyType, TypeOfAny, CallableType, NoneType, TypeVarDef, TypeVarType,
     Overloaded, UnionType, FunctionLike, get_proper_type
 )
 from mypy.typeops import make_simplified_union, map_type_from_supertype
@@ -257,6 +258,15 @@ def _get_decorator_optional_bool_argument(
         return default
 
 
+def _make_var(n: str, i: Instance, fullname: Optional[str] = None, is_classvar: bool = False) -> Var:
+    """Create a `Var` for a symbol table."""
+    res = Var(n, i)
+    res.info = i.type
+    res.is_classvar = is_classvar
+    res._fullname = fullname
+    return res
+
+
 def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
                               auto_attribs_default: Optional[bool] = False,
                               frozen_default: bool = False) -> None:
@@ -319,6 +329,66 @@ def attr_class_maker_callback(ctx: 'mypy.plugin.ClassDefContext',
     if frozen:
         _make_frozen(ctx, attributes)
 
+    # Set up the `__attrs_attrs__` class variable.
+    attr_symbol_tn = lookup_fully_qualified("attr.Attribute",
+        ctx.api.modules, raise_on_missing=False)
+
+    if attr_symbol_tn is None or attr_symbol_tn.node is None:
+        # Not expected to happen since the absence of attrs is handled elsewhere,
+        # but we check just in case. Without this we cannot generate
+        # the `__attrs_attrs__` variable, so we don't.
+        return
+
+    attr_type = attr_symbol_tn.node
+
+    tuple_type: TypeInfo = lookup_fully_qualified("builtins.tuple",
+        ctx.api.modules,
+        raise_on_missing=True).node
+
+    # The `__attrs_attrs__` field is a special namedtuple, unique for each
+    # class, containing all gathered attributes.
+    # It's a namedtuple so attributes can be accessed by name:
+    # `X.__attrs_attrs__.a` points to the attribute for field `X.a`
+    # Both the namedtuple class and instance of it are generated at runtime
+    # by attrs, so we replicate that behavior here.
+
+    sym_table = SymbolTable(
+        {
+            a.name: SymbolTableNode(
+                MDEF,
+                _make_var(
+                    a.name,
+                    Instance(attr_type,
+                        # We want to support future attrs versions with a 2-param Attribute
+                        [Instance(ctx.cls.info, []), info[a.name].type]
+                        if len(attr_type.type_vars) == 2 else
+                        [info[a.name].type]),
+                ),
+            )
+            for a in attributes
+        },
+    )
+    cd = ClassDef(f"{info.name}Attributes", Block([]))
+    cd.fullname = f"builtins.{info.name}Attributes"  # To match the runtime semantics
+    ti = TypeInfo(
+        sym_table,
+        cd,
+        "builtins",
+    )
+    ti.is_named_tuple = True
+    ti.mro = [ti, tuple_type]
+    ti.type_vars = []
+    attributes_type = Instance(
+        ti,
+        [],
+    )
+    fn = ctx.cls.fullname + ".__attrs_attrs__"
+    ctx.cls.info.names["__attrs_attrs__"] = SymbolTableNode(
+        MDEF,
+        _make_var("__attrs_attrs__", attributes_type, fullname=fn, is_classvar=True),
+        plugin_generated=True
+    )
+
 
 def _get_frozen(ctx: 'mypy.plugin.ClassDefContext', frozen_default: bool) -> bool:
     """Return whether this class is frozen."""
@@ -732,3 +802,21 @@ def add_method(self,
         """
         self_type = self_type if self_type is not None else self.self_type
         add_method(self.ctx, method_name, args, ret_type, self_type, tvd)
+
+
+def adjust_fields(fc: FunctionContext) -> Type:
+    """Generate the proper return value for an `attr.fields` invocation."""
+    # We fish out the attrs class out of the return type.
+    try:
+        attrs_class: TypeInfo = fc.arg_types[0][0].ret_type.type
+    except Exception:
+        fc.api.fail("Argument is not an attrs class", fc.context)
+        return AnyType(TypeOfAny.from_error)
+
+    # We've placed the appropriate instance on the type already, when we analyzed it.
+    stn = attrs_class.get("__attrs_attrs__")
+    if stn is not None:
+        return stn.type
+
+    # Should never happen.
+    return AnyType(TypeOfAny.from_error)

mypy/plugins/default.py

@@ -22,14 +22,16 @@ class DefaultPlugin(Plugin):
 
     def get_function_hook(self, fullname: str
                           ) -> Optional[Callable[[FunctionContext], Type]]:
-        from mypy.plugins import ctypes
+        from mypy.plugins import attrs, ctypes
 
         if fullname == 'contextlib.contextmanager':
             return contextmanager_callback
         elif fullname == 'builtins.open' and self.python_version[0] == 3:
             return open_callback
         elif fullname == 'ctypes.Array':
             return ctypes.array_constructor_callback
+        elif fullname == 'attr.fields':
+            return attrs.adjust_fields
         return None
 
     def get_method_signature_hook(self, fullname: str