Improve the collections_abc overlay.

* Handle typing containers better.
* make collections_abc.Set an alias for typing.AbstractSet: https://github.com/python/typeshed/blo[]fc7768c4df6f713e81a0078e72b9805679c25f/stdlib/_collections_abc.pyi#L4

PiperOrigin-RevId: 446856323

Author: rchen152

pytype/overlays/collections_overlay.py

@@ -34,4 +34,4 @@ class ABCOverlay(typing_overlay.Redirect):
   """A custom overlay for the 'collections.abc' module."""
 
   def __init__(self, ctx):
-    super().__init__("collections.abc", {}, ctx)
+    super().__init__("collections.abc", {"Set": "typing.AbstractSet"}, ctx)

pytype/overlays/typing_overlay.py

@@ -23,6 +23,10 @@
 Param = overlay_utils.Param
 
 
+def _is_typing_container(cls: pytd.Class):
+  return pytd.IsContainer(cls) and cls.template
+
+
 class TypingOverlay(overlay.Overlay):
   """A representation of the 'typing' module that allows custom overlays.
 
@@ -39,7 +43,7 @@ def __init__(self, ctx):
     ast = ctx.loader.typing
     for cls in ast.classes:
       _, name = cls.name.rsplit(".", 1)
-      if name not in member_map and pytd.IsContainer(cls) and cls.template:
+      if name not in member_map and _is_typing_container(cls):
         member_map[name] = (overlay.build(name, TypingContainer), None)
     super().__init__(ctx, "typing", member_map, ast)
 
@@ -70,7 +74,7 @@ def __init__(self, module_name, aliases, ctx):
     for pyval in ast.aliases + ast.classes + ast.constants + ast.functions:
       # Any public members that are not explicitly implemented are unsupported.
       _, name = pyval.name.rsplit(".", 1)
-      if name.startswith("_"):
+      if name.startswith("_") or name in member_map:
         continue
       if name in typing_overlay:
         member_map[name] = typing_overlay[name][0]
@@ -84,6 +88,9 @@ def __init__(self, module_name, aliases, ctx):
 def _build(name):
   def resolve(ctx):
     ast = ctx.loader.typing
+    pytd_val = ast.Lookup(name)
+    if isinstance(pytd_val, pytd.Class) and _is_typing_container(pytd_val):
+      return TypingContainer(name.rsplit(".", 1)[-1], ctx)
     pytd_type = pytd.ToType(ast.Lookup(name), True, True, True)
     return ctx.convert.constant_to_value(pytd_type)
   return resolve

pytype/tests/test_collections_abc.py

@@ -53,6 +53,21 @@ def f() -> Callable[[], float]: ...
         x: float
       """)
 
+  def test_generator(self):
+    self.Check("""
+      from collections.abc import Generator
+      def f() -> Generator[int, None, None]:
+        yield 0
+    """)
+
+  def test_set(self):
+    # collections.abc.Set is an alias for typing.AbstractSet.
+    self.Check("""
+      from collections.abc import Set
+      def f() -> Set[int]:
+        return frozenset([0])
+    """)
+
 
 if __name__ == "__main__":
   test_base.main()