Add a stub and overlay for the 'attrs' module.

The attrs library introduced "next-generation" APIs a couple years ago and
recently introduced a new namespace, attrs, which contains only the next-gen
APIs (unlike the old one, attr, which also contains the historical APIs).

The new attrs/__init__.pytd stub is copied from the attrs GitHub repo (from a
commit back in December, in order to be compatible with our attr/ stubs).

PiperOrigin-RevId: 454257324

pytype/overlays/attr_overlay.py

@@ -2,7 +2,7 @@
 
 import enum
 import logging
-from typing import Any, ClassVar, Dict, Optional, Tuple, Union
+from typing import Any, Callable, ClassVar, Dict, Optional, Tuple, TypeVar, Union
 
 from pytype.abstract import abstract
 from pytype.abstract import abstract_utils
@@ -18,6 +18,8 @@
 Param = overlay_utils.Param
 Attribute = classgen.Attribute
 
+_TBaseValue = TypeVar("_TBaseValue", bound=abstract.BaseValue)
+
 
 class TypeSource(enum.Enum):
   """Source of an attrib's `typ` property."""
@@ -26,34 +28,61 @@ class TypeSource(enum.Enum):
   CONVERTER = 3
 
 
-class AttrOverlay(overlay.Overlay):
-  """A custom overlay for the 'attr' module."""
+class _AttrOverlayBase(overlay.Overlay):
+  """Base class for the attr and attrs modules, containing common attributes."""
+
+  _MODULE_NAME: str
 
   def __init__(self, ctx):
     member_map = {
-        "attrs": Attrs.make,
-        "attrib": Attrib.make,
-        "s": Attrs.make,
-        "dataclass": Attrs.make_dataclass,
-        "ib": Attrib.make,
-
         # Attr's next-gen APIs
         # See https://www.attrs.org/en/stable/api.html#next-gen
-
-        # They accept (almost) all the same arguments as the previous APIs.
-        # Strictly speaking these only exist in Python 3.6 and up, but
-        # Python 3.6 is the minimum version pytype are supports now, so not
-        # bothering to guard these definitions.
-        "define": AttrsNextGenDefine.make,
+        "define": self._build(AttrsNextGenDefine.make),
         # These (may) have different default arguments than 'define' for the
         # "frozen" arguments, but the overlay doesn't look at that yet.
-        "mutable": AttrsNextGenDefine.make,
-        "frozen": AttrsNextGenDefine.make,
-        "field": Attrib.make,
+        "mutable": self._build(AttrsNextGenDefine.make),
+        "frozen": self._build(AttrsNextGenDefine.make),
+        "field": self._build(Attrib.make),
     }
+    ast = ctx.loader.import_name(self._MODULE_NAME)
+    super().__init__(ctx, self._MODULE_NAME, member_map, ast)
 
-    ast = ctx.loader.import_name("attr")
-    super().__init__(ctx, "attr", member_map, ast)
+  @classmethod
+  def _build(
+      cls, make_fn: Callable[[Any, str], _TBaseValue],
+  ) -> Callable[[Any], _TBaseValue]:
+    return lambda ctx: make_fn(ctx, cls._MODULE_NAME)
+
+
+class AttrOverlay(_AttrOverlayBase):
+  """A custom overlay for the 'attr' module.
+
+  'attr' is the historical namespace for the attrs library, containing both
+  the old APIs (attr.s, attr.ib, etc.) and the next-generation ones
+  (attr.define, attr.field, etc.)
+  """
+
+  _MODULE_NAME = "attr"
+
+  def __init__(self, ctx):
+    super().__init__(ctx)
+    self._member_map.update({
+        "attrs": Attrs.make,
+        "attrib": self._build(Attrib.make),
+        "s": Attrs.make,
+        "dataclass": Attrs.make_dataclass,
+        "ib": self._build(Attrib.make),
+    })
+
+
+class AttrsOverlay(_AttrOverlayBase):
+  """A custom overlay for the 'attrs' module.
+
+  'attrs' is the new namespace for the attrs library's next-generation APIs
+  (attrs.define, attrs.field, etc.)
+  """
+
+  _MODULE_NAME = "attrs"
 
 
 class _NoChange():
@@ -227,9 +256,9 @@ class AttrsNextGenDefine(Attrs):
   }
 
   @classmethod
-  def make(cls, ctx):
+  def make(cls, ctx, module):
     # Bypass Attrs's make; go straight to its superclass.
-    return super(Attrs, cls).make("define", ctx, "attr")  # pylint: disable=bad-super-call
+    return super(Attrs, cls).make("define", ctx, module)  # pylint: disable=bad-super-call
 
   def _handle_auto_attribs(self, auto_attribs, local_ops, cls_name):
     if auto_attribs is not None:
@@ -320,11 +349,11 @@ def from_metadata(cls, ctx, node, typ, metadata):
 
 
 class Attrib(classgen.FieldConstructor):
-  """Implements attr.ib."""
+  """Implements attr.ib/attrs.field."""
 
   @classmethod
-  def make(cls, ctx):
-    return super().make("ib", ctx, "attr")
+  def make(cls, ctx, module):
+    return super().make("ib" if module == "attr" else "field", ctx, module)
 
   def _match_and_discard_args(self, node, funcb, args):
     """Discard invalid args so that we can still construct an attrib."""

pytype/overlays/overlay_dict.py

@@ -30,6 +30,7 @@
     "abc": abc_overlay.ABCOverlay,
     "asyncio": asyncio_types_overlay.AsyncioOverlay,
     "attr": attr_overlay.AttrOverlay,
+    "attrs": attr_overlay.AttrsOverlay,
     "chex": chex_overlay.ChexOverlay,
     "collections": collections_overlay.CollectionsOverlay,
     "collections.abc": collections_overlay.ABCOverlay,

pytype/stubs/CMakeLists.txt

@@ -12,6 +12,7 @@ filegroup(
     builtins/attr/filters.pytd
     builtins/attr/setters.pytd
     builtins/attr/validators.pytd
+    builtins/attrs/__init__.pytd
     builtins/builtins.pytd
     builtins/mypy_extensions.pytd
     builtins/numpy/__init__.pytd

pytype/stubs/builtins/attrs/__init__.pytd

@@ -0,0 +1,63 @@
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+)
+
+# Because we need to type our own stuff, we have to make everything from
+# attr explicitly public too.
+from attr import __author__ as __author__
+from attr import __copyright__ as __copyright__
+from attr import __description__ as __description__
+from attr import __email__ as __email__
+from attr import __license__ as __license__
+from attr import __title__ as __title__
+from attr import __url__ as __url__
+from attr import __version__ as __version__
+from attr import __version_info__ as __version_info__
+from attr import _FilterType
+from attr import assoc as assoc
+from attr import Attribute as Attribute
+from attr import define as define
+from attr import evolve as evolve
+from attr import Factory as Factory
+from attr import exceptions as exceptions
+from attr import field as field
+from attr import fields as fields
+from attr import fields_dict as fields_dict
+from attr import frozen as frozen
+from attr import has as has
+from attr import make_class as make_class
+from attr import mutable as mutable
+from attr import NOTHING as NOTHING
+from attr import resolve_types as resolve_types
+from attr import setters as setters
+from attr import validate as validate
+from attr import validators as validators
+
+# TODO: see definition of attr.asdict/astuple
+def asdict(
+    inst: Any,
+    recurse: bool = ...,
+    filter: Optional[_FilterType[Any]] = ...,
+    dict_factory: Type[Mapping[Any, Any]] = ...,
+    retain_collection_types: bool = ...,
+    value_serializer: Optional[
+        Callable[[type, Attribute[Any], Any], Any]
+    ] = ...,
+    tuple_keys: bool = ...,
+) -> Dict[str, Any]: ...
+
+# TODO: add support for returning NamedTuple from the mypy plugin
+def astuple(
+    inst: Any,
+    recurse: bool = ...,
+    filter: Optional[_FilterType[Any]] = ...,
+    tuple_factory: Type[Sequence[Any]] = ...,
+    retain_collection_types: bool = ...,
+) -> Tuple[Any, ...]: ...

pytype/tests/test_attr2.py

@@ -954,6 +954,43 @@ class Foo:
         def __init__(self, y: str = ...) -> None: ...
     """)
 
+  def test_attrs_namespace(self):
+    ty = self.Infer("""
+      import attrs
+      @attrs.define
+      class Foo:
+        x: int
+      @attrs.mutable
+      class Bar:
+        x: int
+      @attrs.frozen
+      class Baz:
+        x: int
+      @attrs.define
+      class Qux:
+        x: int = attrs.field(init=False)
+    """)
+    self.assertTypesMatchPytd(ty, """
+      import attr
+      import attrs
+      @attr.s
+      class Foo:
+        x: int
+        def __init__(self, x: int) -> None: ...
+      @attr.s
+      class Bar:
+        x: int
+        def __init__(self, x: int) -> None: ...
+      @attr.s
+      class Baz:
+        x: int
+        def __init__(self, x: int) -> None: ...
+      @attr.s
+      class Qux:
+        x: int
+        def __init__(self) -> None: ...
+    """)
+
 
 class TestPyiAttrs(test_base.BaseTest):
   """Tests for @attr.s in pyi files."""