Add signature for attr.evolve

mypy/plugins/attrs.py

@@ -6,8 +6,10 @@
 from typing_extensions import Final, Literal
 
 import mypy.plugin # To avoid circular imports.
+from mypy.checker import TypeChecker
 from mypy.errorcodes import LITERAL_REQ
 from mypy.exprtotype import TypeTranslationError, expr_to_unanalyzed_type
+from mypy.messages import format_type_bare
 from mypy.nodes import (
  ARG_NAMED,
  ARG_NAMED_OPT,
@@ -77,6 +79,7 @@
 SELF_TVAR_NAME: Final = "_AT"
 MAGIC_ATTR_NAME: Final = "__attrs_attrs__"
 MAGIC_ATTR_CLS_NAME_TEMPLATE: Final = "__{}_AttrsAttributes__" # The tuple subclass pattern.
+ATTRS_INIT_NAME: Final = "__attrs_init__"
 
 
 class Converter:
@@ -330,7 +333,7 @@ def attr_class_maker_callback(
 
  adder = MethodAdder(ctx)
  # If __init__ is not being generated, attrs still generates it as __attrs_init__ instead.
- _add_init(ctx, attributes, adder, "__init__" if init else "__attrs_init__")
+ _add_init(ctx, attributes, adder, "__init__" if init else ATTRS_INIT_NAME)
  if order:
  _add_order(ctx, adder)
  if frozen:
@@ -888,3 +891,64 @@ def add_method(
  """
  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 _get_attrs_init_type(typ: Type) -> CallableType | None:
+ """
+ If `typ` refers to an attrs class, gets the type of its initializer method.
+ """
+ typ = get_proper_type(typ)
+ if not isinstance(typ, Instance):
+ return None
+ magic_attr = typ.type.get(MAGIC_ATTR_NAME)
+ if magic_attr is None or not magic_attr.plugin_generated:
+ return None
+ init_method = typ.type.get_method("__init__") or typ.type.get_method(ATTRS_INIT_NAME)
+ if not isinstance(init_method, FuncDef) or not isinstance(init_method.type, CallableType):
+ return None
+ return init_method.type
+
+
+def evolve_function_sig_callback(ctx: mypy.plugin.FunctionSigContext) -> CallableType:
+ """
+ Generates a signature for the 'attr.evolve' function that's specific to the call site
+ and dependent on the type of the first argument.
+ """
+ if len(ctx.args) != 2:
+ # Ideally the name and context should be callee's, but we don't have it in FunctionSigContext.
+ ctx.api.fail(f'"{ctx.default_signature.name}" has unexpected type annotation', ctx.context)
+ return ctx.default_signature
+
+ if len(ctx.args[0]) != 1:
+ return ctx.default_signature # leave it to the type checker to complain
+
+ inst_arg = ctx.args[0][0]
+
+ # <hack>
+ assert isinstance(ctx.api, TypeChecker)
+ inst_type = ctx.api.expr_checker.accept(inst_arg)
+ # </hack>
+
+ inst_type = get_proper_type(inst_type)
+ if isinstance(inst_type, AnyType):
+ return ctx.default_signature
+ inst_type_str = format_type_bare(inst_type)
+
+ attrs_init_type = _get_attrs_init_type(inst_type)
+ if not attrs_init_type:
+ ctx.api.fail(
+ f'Argument 1 to "evolve" has incompatible type "{inst_type_str}"; expected an attrs class',
+ ctx.context,
+ )
+ return ctx.default_signature
+
+ # AttrClass.__init__ has the following signature (or similar, if having kw-only & defaults):
+ # def __init__(self, attr1: Type1, attr2: Type2) -> None:
+ # We want to generate a signature for evolve that looks like this:
+ # def evolve(inst: AttrClass, *, attr1: Type1 = ..., attr2: Type2 = ...) -> AttrClass:
+ return attrs_init_type.copy_modified(
+ arg_names=["inst"] + attrs_init_type.arg_names[1:],
+ arg_kinds=[ARG_POS] + [ARG_NAMED_OPT for _ in attrs_init_type.arg_kinds[1:]],
+ ret_type=inst_type,
+ name=f"{ctx.default_signature.name} of {inst_type_str}",
+ )

mypy/plugins/default.py

@@ -10,6 +10,7 @@
  AttributeContext,
  ClassDefContext,
  FunctionContext,
+ FunctionSigContext,
  MethodContext,
  MethodSigContext,
  Plugin,
@@ -46,6 +47,15 @@ def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type]
  return singledispatch.create_singledispatch_function_callback
  return None
 
+ def get_function_signature_hook(
+ self, fullname: str
+ ) -> Callable[[FunctionSigContext], FunctionLike] | None:
+ from mypy.plugins import attrs
+
+ if fullname in ("attr.evolve", "attrs.evolve", "attr.assoc", "attrs.assoc"):
+ return attrs.evolve_function_sig_callback
+ return None
+
  def get_method_signature_hook(
  self, fullname: str
  ) -> Callable[[MethodSigContext], FunctionLike] | None:

mypy/types.py

@@ -1775,7 +1775,7 @@ def copy_modified(
  self: CT,
  arg_types: Bogus[Sequence[Type]] = _dummy,
  arg_kinds: Bogus[list[ArgKind]] = _dummy,
- arg_names: Bogus[list[str | None]] = _dummy,
+ arg_names: Bogus[Sequence[str | None]] = _dummy,
  ret_type: Bogus[Type] = _dummy,
  fallback: Bogus[Instance] = _dummy,
  name: Bogus[str | None] = _dummy,

test-data/unit/check-attr.test

@@ -1867,3 +1867,81 @@ D(1, "").a = 2 # E: Cannot assign to final attribute "a"
 D(1, "").b = "2" # E: Cannot assign to final attribute "b"
 
 [builtins fixtures/property.pyi]
+
+[case testEvolve]
+import attr
+
+class Base:
+ pass
+
+class Derived(Base):
+ pass
+
+class Other:
+ pass
+
+@attr.s(auto_attribs=True)
+class C:
+ name: str
+ b: Base
+
+c = C(name='foo', b=Derived())
+c = attr.evolve(c)
+c = attr.evolve(c, name='foo')
+c = attr.evolve(c, 'foo') # E: Too many positional arguments for "evolve" of "C"
+c = attr.evolve(c, b=Derived())
+c = attr.evolve(c, b=Base())
+c = attr.evolve(c, b=Other()) # E: Argument "b" to "evolve" of "C" has incompatible type "Other"; expected "Base"
+c = attr.evolve(c, name=42) # E: Argument "name" to "evolve" of "C" has incompatible type "int"; expected "str"
+c = attr.evolve(c, foobar=42) # E: Unexpected keyword argument "foobar" for "evolve" of "C"
+
+# test passing instance as 'inst' kw
+c = attr.evolve(inst=c, name='foo')
+c = attr.evolve(not_inst=c, name='foo') # E: Missing positional argument "inst" in call to "evolve"
+
+# test determining type of first argument's expression from something that's not NameExpr
+def f() -> C:
+ return c
+
+c = attr.evolve(f(), name='foo')
+
+[builtins fixtures/attr.pyi]
+
+[case testEvolveFromNonAttrs]
+import attr
+
+attr.evolve(42, name='foo') # E: Argument 1 to "evolve" has incompatible type "int"; expected an attrs class
+attr.evolve(None, name='foo') # E: Argument 1 to "evolve" has incompatible type "None"; expected an attrs class
+[case testEvolveFromAny]
+from typing import Any
+import attr
+
+any: Any = 42
+ret = attr.evolve(any, name='foo')
+reveal_type(ret) # N: Revealed type is "Any"
+
+[typing fixtures/typing-medium.pyi]
+
+[case testEvolveVariants]
+from typing import Any
+import attr
+import attrs
+
+
+@attr.s(auto_attribs=True)
+class C:
+ name: str
+
+c = C(name='foo')
+
+c = attr.assoc(c, name='test')
+c = attr.assoc(c, name=42) # E: Argument "name" to "assoc" of "C" has incompatible type "int"; expected "str"
+
+c = attrs.evolve(c, name='test')
+c = attrs.evolve(c, name=42) # E: Argument "name" to "evolve" of "C" has incompatible type "int"; expected "str"
+
+c = attrs.assoc(c, name='test')
+c = attrs.assoc(c, name=42) # E: Argument "name" to "assoc" of "C" has incompatible type "int"; expected "str"
+
+[builtins fixtures/attr.pyi]
+[typing fixtures/typing-medium.pyi]

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

@@ -244,3 +244,6 @@ def field(
  order: Optional[bool] = ...,
  on_setattr: Optional[object] = ...,
 ) -> Any: ...
+
+def evolve(inst: _T, **changes: Any) -> _T: ...
+def assoc(inst: _T, **changes: Any) -> _T: ...

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

@@ -126,3 +126,6 @@ def field(
  order: Optional[bool] = ...,
  on_setattr: Optional[object] = ...,
 ) -> Any: ...
+
+def evolve(inst: _T, **changes: Any) -> _T: ...
+def assoc(inst: _T, **changes: Any) -> _T: ...