Google sync

pytype/overlays/enum_overlay.py

@@ -44,9 +44,8 @@
 
 
 # These members have been added in Python 3.11 and are not yet supported.
-_unsupported = ("StrEnum", "ReprEnum", "EnumCheck", "FlagBoundary", "verify",
-                "property", "member", "nonmember", "global_enum",
-                "show_flag_values")
+_unsupported = ("ReprEnum", "EnumCheck", "FlagBoundary", "verify", "property",
+                "member", "nonmember", "global_enum", "show_flag_values")
 
 
 class EnumOverlay(overlay.Overlay):
@@ -59,6 +58,7 @@ def __init__(self, ctx):
           "EnumMeta": EnumMeta,
           "EnumType": EnumMeta,
           "IntEnum": overlay.add_name("IntEnum", EnumBuilder),
+          "StrEnum": overlay.add_name("StrEnum", EnumBuilder),
           **{name: overlay.add_name(name, overlay_utils.not_supported_yet)
              for name in _unsupported},
       }

pytype/pyc/opcodes.py

@@ -1122,20 +1122,17 @@ def _make_opcode_list(offset_to_op, python_version):
   op_items = sorted(offset_to_op.items())
   for i, (off, op) in enumerate(op_items):
     index += 1
-    if python_version == (3, 11):
-      if (isinstance(op, JUMP_BACKWARD) and
-          i + 1 < len(op_items) and
-          isinstance(op_items[i + 1][1], END_ASYNC_FOR)):
-        # In 3.11 `async for` is compiled into an infinite loop, relying on the
-        # exception handler to break out. This causes the block graph to be
-        # pruned abruptly, so we need to remove the loop opcode.
-        index -= 1
-        continue
-      elif (isinstance(op, JUMP_BACKWARD_NO_INTERRUPT) and
-            isinstance(offset_to_op[op.argval], SEND)):
-        # Likewise, `await` is compiled into an infinite loop which we remove.
-        index -= 1
-        continue
+    if (python_version == (3, 11) and isinstance(op, JUMP_BACKWARD) and
+        i + 1 < len(op_items) and
+        isinstance(op_items[i + 1][1], END_ASYNC_FOR)):
+      # In 3.11 `async for` is compiled into an infinite loop, relying on the
+      # exception handler to break out. This causes the block graph to be
+      # pruned abruptly, so we need to remove the loop opcode.
+      #  We map the offset to the index of the next opcode so that jumps to
+      # `op` are redirected correctly.
+      offset_to_index[off] = index
+      index -= 1
+      continue
     op.index = index
     offset_to_index[off] = index
     if prev_op:
@@ -1151,10 +1148,7 @@ def _add_jump_targets(ops, offset_to_index):
   """Map the target of jump instructions to the opcode they jump to."""
   for op in ops:
     op = cast(OpcodeWithArg, op)
-    if isinstance(op, SEND):
-      # This has a target in the bytecode, but is not a jump
-      op.target = None
-    elif op.target:
+    if op.target:
       # We have already set op.target, we need to fill in its index in op.arg
       op.arg = op.argval = op.target.index
     elif op.has_known_jump():

pytype/pyi/parse_pyi.py

@@ -18,7 +18,7 @@
   module_name = module_utils.path_to_module_name(filename)
 
   try:
-    out, _ = parser.parse_pyi(src, filename, module_name, debug_mode=True)
+    out = parser.parse_pyi(src, filename, module_name, debug_mode=True)
   except _ParseError as e:
     print(e)
     sys.exit(1)

pytype/pytd/pytd_visitors.py

@@ -157,7 +157,7 @@ def _MaybeNewName(self, name):
     if name == self._old[:-1]:
       return self._module_name
     before, match, after = name.partition(self._old)
-    if match and not before and "." not in after:
+    if match and not before:
       return self._new + after
     else:
       return name

pytype/stubs/builtins/builtins.pytd

@@ -694,8 +694,9 @@ class set(Set[_T]):
         self = set[Union[_T, _T2]]
     @overload
     def union(self, other: Iterable[_T2]) -> set[_T | _T2]: ...
-    @overload
-    def union(self, other1: Iterable[_T2], other2: Iterable[_T3]) -> set[_T | _T2 | _T3]: ...
+    # TODO(b/304591130): Re-enable this overload.
+    # @overload
+    # def union(self, other1: Iterable[_T2], other2: Iterable[_T3]) -> set[_T | _T2 | _T3]: ...
     @overload
     def union(self, *args: Iterable[Any]) -> set[Any]: ...
     def intersection(self, *args: Iterable[Any]) -> set[_T]: ...
@@ -720,8 +721,9 @@ class frozenset(FrozenSet[_T]):
     def issuperset(self, y: Iterable) -> bool: ...
     @overload
     def union(self, other: Iterable[_T2]) -> frozenset[_T | _T2]: ...
-    @overload
-    def union(self, other1: Iterable[_T2], other2: Iterable[_T3]) -> frozenset[_T | _T2 | _T3]: ...
+    # TODO(b/304591130): Re-enable this overload.
+    # @overload
+    # def union(self, other1: Iterable[_T2], other2: Iterable[_T3]) -> frozenset[_T | _T2 | _T3]: ...
     @overload
     def union(self, *args: Iterable[Any]) -> frozenset[Any]: ...
     def intersection(self, *args: Iterable[Any]) -> frozenset[_T]: ...

pytype/stubs/stdlib/enum.pytd

@@ -66,7 +66,12 @@ class IntFlag(int, Flag):
 
 # 3.11: Not supported yet.
 EnumType = EnumMeta
-class StrEnum: ...
+
+class StrEnum(str, Enum):
+  # Workaround for b/201603421.
+  @classmethod
+  def __iter__(cls: Type[_T]) -> Iterator[_T]: ...
+
 class ReprEnum: ...
 class EnumCheck: ...
 class FlagBoundary: ...

pytype/tests/CMakeLists.txt

@@ -531,6 +531,15 @@ py_test(
     .test_base
 )
 
+py_test(
+  NAME
+    test_flow3
+  SRCS
+    test_flow3.py
+  DEPS
+    .test_base
+)
+
 py_test(
   NAME
     test_methods1
@@ -1255,3 +1264,12 @@ py_test(
   DEPS
     .test_base
 )
+
+py_test(
+  NAME
+    test_typing_self
+  SRCS
+    test_typing_self.py
+  DEPS
+    .test_base
+)

pytype/tests/test_enums.py

@@ -924,6 +924,15 @@ class IF(enum.IntFlag):
         A = 1
     """)
 
+  def test_strenum(self):
+    self.Check("""
+      import enum
+      class MyEnum(enum.StrEnum):
+        A = 'A'
+      for x in MyEnum:
+        assert_type(x, MyEnum)
+    """)
+
   def test_unique_enum_in_dict(self):
     # Regression test for a recursion error in matcher.py
     self.assertNoCrash(self.Check, """
@@ -1358,7 +1367,7 @@ class M(enum.Enum):
   def test_not_supported_yet(self):
     self.CheckWithErrors("""
       import enum
-      enum.StrEnum  # not-supported-yet
+      enum.ReprEnum  # not-supported-yet
     """)
 
 

pytype/tests/test_flow3.py

@@ -0,0 +1,145 @@
+"""Tests for control flow cases that involve the exception table in 3.11+.
+
+Python 3.11 changed the way exceptions and some other control structures were
+compiled, and in particular some of them require examining the exception table
+as well as the bytecode.
+"""
+
+from pytype.tests import test_base
+
+
+class TestPy311(test_base.BaseTest):
+  """Tests for python 3.11 support."""
+
+  def test_context_manager(self):
+    self.Check("""
+      class A:
+        def __enter__(self):
+          pass
+        def __exit__(self, a, b, c):
+          pass
+
+      lock = A()
+
+      def f() -> str:
+        path = ''
+        with lock:
+          try:
+            pass
+          except:
+            pass
+          return path
+    """)
+
+  def test_exception_type(self):
+    self.Check("""
+      class FooError(Exception):
+        pass
+      try:
+        raise FooError()
+      except FooError as e:
+        assert_type(e, FooError)
+    """)
+
+  def test_try_with(self):
+    self.Check("""
+      def f(obj, x):
+        try:
+          with __any_object__:
+            obj.get(x)
+        except:
+          pass
+    """)
+
+  def test_try_if_with(self):
+    self.Check("""
+      from typing import Any
+      import os
+      pytz: Any
+      def f():
+        tz_env = os.environ.get('TZ')
+        try:
+          if tz_env == 'localtime':
+            with open('localtime') as localtime:
+              return pytz.tzfile.build_tzinfo('', localtime)
+        except IOError:
+          return pytz.UTC
+    """)
+
+  def test_try_finally(self):
+    self.Check("""
+      import tempfile
+      dir_ = None
+      def f():
+        global dir_
+        try:
+          if dir_:
+            return dir_
+          dir_ = tempfile.mkdtemp()
+        finally:
+          print(dir_)
+    """)
+
+  def test_nested_try_in_for(self):
+    self.Check("""
+      def f(x):
+        for i in x:
+          fd = __any_object__
+          try:
+            try:
+              if __random__:
+                return True
+            except ValueError:
+              continue
+          finally:
+            fd.close()
+    """)
+
+  def test_while_and_nested_try(self):
+    self.Check("""
+      def f(p):
+        try:
+          while __random__:
+            try:
+              return p.communicate()
+            except KeyboardInterrupt:
+              pass
+        finally:
+          pass
+    """)
+
+  def test_while_and_nested_try_2(self):
+    self.Check("""
+      def f():
+        i = j = 0
+        while True:
+          try:
+            try:
+              i += 1
+            finally:
+              j += 1
+          except:
+            break
+        return
+    """)
+
+  def test_while_and_nested_try_3(self):
+    self.Check("""
+      import os
+
+      def RmDirs(dir_name):
+        try:
+          parent_directory = os.path.dirname(dir_name)
+          while parent_directory:
+            try:
+              os.rmdir(parent_directory)
+            except OSError as err:
+              pass
+            parent_directory = os.path.dirname(parent_directory)
+        except OSError as err:
+          pass
+    """)
+
+
+if __name__ == "__main__":
+  test_base.main()