Introduce Y052: Disallow assignments to constant values in global or …

…class namespaces where the assignments don't have type annotations (#339)

Following #326, we now allow constructs such as these in a stub file:

```python
x = 1
y = "foo"
z = b"bar"
```

It's good that we now allow variables to have default values. However,
the above constructs are problematic in a stub, as we're leaving it up
to the type checker to make inferences about the types of these
variables, which can have unpredictable consequences. For example, the
`x` variable could be inferred to be of type `int`, `Final[int]`,
`Literal[1]`, or `Final[Literal[1]]`. In a class namespace, the problem
is even worse -- the `eggs` variable in the following snippet could
reasonably be inferred as being of type `str`, `ClassVar[str]`,
`Final[str]`, `Literal["EGGS"]`, `ClassVar[Literal["EGGS"]]`, or
`Final[Literal["EGGS"]]`:

```python
class Whatever:
    eggs = "EGGS"
```

This PR introduces Y052, enforcing type annotations for assignments to
simple constants.

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 ## Unreleased
 
+New error codes:
+* Y052: Disallow default values in global or class namespaces where the
+  assignment does not have a type annotation. Stubs should be explicit about
+  the type of all variables in the stub; without type annotations, the type
+  checker is forced to make inferences, which may have unpredictable
+  consequences. Enum members are excluded from this check, as are various
+  special assignments such as `__all__` and `__match_args__`.
+
+Other enhancements:
 * Disallow numeric default values where `len(str(default)) > 7`. If a function
   has a default value where the string representation is greater than 7
   characters, it is likely to be an implementation detail or a constant that

README.md

@@ -83,7 +83,8 @@ currently emitted:
 | Y048 | Function bodies should contain exactly one statement. (Note that if a function body includes a docstring, the docstring counts as a "statement".)
 | Y049 | A private `TypedDict` should be used at least once in the file in which it is defined.
 | Y050 | Prefer `typing_extensions.Never` over `typing.NoReturn` for argument annotations. This is a purely stylistic choice in the name of readability.
-| Y051 | Y051 detect redundant unions between `Literal` types and builtin supertypes. For example, `Literal[5]` is redundant in the union `int \| Literal[5]`, and `Literal[True]` is redundant in the union `Literal[True] \| bool`.
+| Y051 | Y051 detects redundant unions between `Literal` types and builtin supertypes. For example, `Literal[5]` is redundant in the union `int \| Literal[5]`, and `Literal[True]` is redundant in the union `Literal[True] \| bool`.
+| Y052 | Y052 disallows assignments to constant values where the assignment does not have a type annotation. For example, `x = 0` in the global namespace is ambiguous in a stub, as there are four different types that could be inferred for the variable `x`: `int`, `Final[int]`, `Literal[0]`, or `Final[Literal[0]]`. Enum members are excluded from this check, as are various special assignments such as `__all__` and `__match_args__`.
 
 Many error codes enforce modern conventions, and some cannot yet be used in
 all cases:

pyi.py

@@ -369,6 +369,7 @@ def _is_object(node: ast.expr | None, name: str, *, from_: Container[str]) -> bo
 )
 _is_Protocol = partial(_is_object, name="Protocol", from_=_TYPING_MODULES)
 _is_NoReturn = partial(_is_object, name="NoReturn", from_=_TYPING_MODULES)
+_is_Final = partial(_is_object, name="Final", from_=_TYPING_MODULES)
 
 
 def _is_object_or_Unused(node: ast.expr | None) -> bool:
@@ -712,8 +713,13 @@ def _analyse_union(members: Sequence[ast.expr]) -> UnionAnalysis:
 )
 
 
-def _is_valid_stub_default(node: ast.expr) -> bool:
-    """Is `node` valid as a default value for a function or method parameter in a stub?"""
+def _is_valid_default_value_with_annotation(node: ast.expr) -> bool:
+    """Is `node` valid as a default value for a function or method parameter in a stub?
+
+    Note that this function is *also* used to determine
+    the validity of default values for ast.AnnAssign nodes.
+    (E.g. `foo: int = 5` is OK, but `foo: TypeVar = TypeVar("foo")` is not.)
+    """
     # `...`, bools, None
     if isinstance(node, (ast.Ellipsis, ast.NameConstant)):
         return True
@@ -769,6 +775,7 @@ def _is_valid_pep_604_union_member(node: ast.expr) -> bool:
 
 
 def _is_valid_pep_604_union(node: ast.expr) -> TypeGuard[ast.BinOp]:
+    """Does `node` represent a valid PEP-604 union (e.g. `int | str`)?"""
     return (
         isinstance(node, ast.BinOp)
         and isinstance(node.op, ast.BitOr)
@@ -780,15 +787,37 @@ def _is_valid_pep_604_union(node: ast.expr) -> TypeGuard[ast.BinOp]:
     )
 
 
-def _is_valid_assignment_value(node: ast.expr) -> bool:
-    """Is `node` valid as the default value for an assignment in a stub?"""
+def _is_valid_default_value_without_annotation(node: ast.expr) -> bool:
+    """Is `node` a valid default for an assignment without an annotation?"""
     return (
-        isinstance(node, (ast.Call, ast.Name, ast.Attribute, ast.Subscript))
+        isinstance(
+            node, (ast.Call, ast.Name, ast.Attribute, ast.Subscript, ast.Ellipsis)
+        )
+        or _is_None(node)
         or _is_valid_pep_604_union(node)
-        or _is_valid_stub_default(node)
     )
 
 
+_KNOWN_ENUM_BASES = frozenset(
+    {"Enum", "Flag", "IntEnum", "IntFlag", "StrEnum", "ReprEnum"}
+)
+
+
+def _is_enum_base(node: ast.expr) -> bool:
+    if isinstance(node, ast.Name):
+        return node.id in _KNOWN_ENUM_BASES
+    return (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id == "enum"
+        and node.attr in _KNOWN_ENUM_BASES
+    )
+
+
+def _is_enum_class(node: ast.ClassDef) -> bool:
+    return any(_is_enum_base(base) for base in node.bases)
+
+
 @dataclass
 class NestingCounter:
     """Class to help the PyiVisitor keep track of internal state"""
@@ -944,6 +973,22 @@ def _check_for_typevarlike_assignments(
             else:
                 self.error(node, Y001.format(cls_name))
 
+    def _check_default_value_without_type_annotation(
+        self, node: ast.Assign, assignment: ast.expr, target_name: str
+    ) -> None:
+        if _is_valid_default_value_without_annotation(assignment):
+            return
+        if _is_valid_default_value_with_annotation(assignment):
+            # Annoying special-casing to exclude enums from Y052
+            if self.in_class.active:
+                assert self.current_class_node is not None
+                if not _is_enum_class(self.current_class_node):
+                    self.error(node, Y052.format(variable=target_name))
+            else:
+                self.error(node, Y052.format(variable=target_name))
+        else:
+            self.error(node, Y015)
+
     def visit_Assign(self, node: ast.Assign) -> None:
         if self.in_function.active:
             # We error for unexpected things within functions separately.
@@ -984,8 +1029,9 @@ def visit_Assign(self, node: ast.Assign) -> None:
 
         if not is_special_assignment:
             self._check_for_type_aliases(node, target, assignment)
-            if not _is_valid_assignment_value(assignment):
-                self.error(node, Y015)
+            self._check_default_value_without_type_annotation(
+                node, assignment, target_name
+            )
 
     def visit_AugAssign(self, node: ast.AugAssign) -> None:
         """Allow `__all__ += ['foo', 'bar']` in a stub file"""
@@ -1120,8 +1166,17 @@ def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
 
         if _is_TypeAlias(node_annotation) and isinstance(node_target, ast.Name):
             self._check_typealias(node=node, alias_name=node_target.id)
+            # Don't bother checking whether
+            # nodes marked as TypeAliases have valid assignment values.
+            # Type checkers will emit errors for those.
+            return
+
+        if _is_Final(node_annotation) and isinstance(
+            node_value, (ast.Attribute, ast.Name)
+        ):
+            return
 
-        if node_value and not _is_valid_assignment_value(node_value):
+        if node_value and not _is_valid_default_value_with_annotation(node_value):
             self.error(node, Y015)
 
     def _check_union_members(self, members: Sequence[ast.expr]) -> None:
@@ -1762,7 +1817,7 @@ def check_arg_default(self, arg: ast.arg, default: ast.expr | None) -> None:
         if default is not None:
             with self.string_literals_allowed.enabled():
                 self.visit(default)
-        if default is not None and not _is_valid_stub_default(default):
+        if default is not None and not _is_valid_default_value_with_annotation(default):
             self.error(default, (Y014 if arg.annotation is None else Y011))
 
     def error(self, node: ast.AST, message: str) -> None:
@@ -1957,3 +2012,4 @@ def parse_options(
     'Y050 Use "typing_extensions.Never" instead of "NoReturn" for argument annotations'
 )
 Y051 = 'Y051 "{literal_subtype}" is redundant in a union with "{builtin_supertype}"'
+Y052 = 'Y052 Need type annotation for "{variable}"'

tests/attribute_annotations.pyi

@@ -1,6 +1,9 @@
 import builtins
+import enum
+import os
 import sys
 import typing
+from enum import Enum, Flag, ReprEnum, StrEnum
 from typing import Final, Final as _Final, TypeAlias
 
 import typing_extensions
@@ -10,17 +13,17 @@ field1: int
 field2: int = ...
 field3 = ...  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
 field4: int = 0
-field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
-field6 = 0
-field7 = b""
-field71 = "foo"
+field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")  # Y052 Need type annotation for "field5"
+field6 = 0  # Y052 Need type annotation for "field6"
+field7 = b""  # Y052 Need type annotation for "field7"
+field71 = "foo"  # Y052 Need type annotation for "field71"
 field72: str = "foo"
-field8 = False
-field81 = -1
+field8 = False  # Y052 Need type annotation for "field8"
+field81 = -1  # Y052 Need type annotation for "field81"
 field82: float = -98.43
-field83 = -42j
-field84 = 5 + 42j
-field85 = -5 - 42j
+field83 = -42j  # Y052 Need type annotation for "field83"
+field84 = 5 + 42j  # Y052 Need type annotation for "field84"
+field85 = -5 - 42j  # Y052 Need type annotation for "field85"
 field9 = None  # Y026 Use typing_extensions.TypeAlias for type aliases, e.g. "field9: TypeAlias = None"
 Field95: TypeAlias = None
 Field96: TypeAlias = int | None
@@ -35,6 +38,9 @@ field15: _Final = True
 field16: typing.Final = "foo"
 field17: typing_extensions.Final = "foo"
 field18: Final = -24j
+field181: Final = field18
+field182: Final = os.pathsep
+field183: Final = None
 
 # We *should* emit Y015 for more complex default values
 field19 = [1, 2, 3]  # Y015 Only simple default values are allowed for assignments
@@ -50,12 +56,12 @@ class Foo:
     field2: int = ...
     field3 = ...  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
     field4: int = 0
-    field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
-    field6 = 0
-    field7 = b""
-    field71 = "foo"
+    field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")  # Y052 Need type annotation for "field5"
+    field6 = 0  # Y052 Need type annotation for "field6"
+    field7 = b""  # Y052 Need type annotation for "field7"
+    field71 = "foo"  # Y052 Need type annotation for "field71"
     field72: str = "foo"
-    field8 = False
+    field8 = False  # Y052 Need type annotation for "field8"
     # Tests for Final
     field9: Final = 1
     field10: Final = "foo"
@@ -65,13 +71,13 @@ class Foo:
     field14: typing.Final = "foo"
     field15: typing_extensions.Final = "foo"
     # Standalone strings used to cause issues
-    field16 = "x"
+    field16 = "x"  # Y052 Need type annotation for "field16"
     if sys.platform == "linux":
-        field17 = "y"
+        field17 = "y"  # Y052 Need type annotation for "field17"
     elif sys.platform == "win32":
-        field18 = "z"
+        field18 = "z"  # Y052 Need type annotation for "field18"
     else:
-        field19 = "w"
+        field19 = "w"  # Y052 Need type annotation for "field19"
 
     field20 = [1, 2, 3]  # Y015 Only simple default values are allowed for assignments
     field21 = (1, 2, 3)  # Y015 Only simple default values are allowed for assignments
@@ -84,3 +90,25 @@ class Foo:
     Field95: TypeAlias = None
     Field96: TypeAlias = int | None
     Field97: TypeAlias = None | typing.SupportsInt | builtins.str
+
+# Enums are excluded from Y052
+class Enum1(Enum):
+    FOO = "foo"
+
+class Enum2(enum.IntEnum):
+    FOO = 1
+
+class Enum3(Flag):
+    FOO = 1
+
+class Enum4(enum.IntFlag):
+    FOO = 1
+
+class Enum5(StrEnum):
+    FOO = "foo"
+
+class Enum6(ReprEnum):
+    FOO = "foo"
+
+class Enum7(enum.Enum):
+    FOO = "foo"

tests/quotes.pyi

@@ -57,8 +57,8 @@ class DocstringAndPass:
     pass  # Y012 Class body must not contain "pass"
 
 # These two shouldn't trigger Y020 -- empty strings can't be "quoted annotations"
-k = ""
-el = r""
+k = ""  # Y052 Need type annotation for "k"
+el = r""  # Y052 Need type annotation for "el"
 
 # The following should also pass,
 # But we can't test for it in CI, because the error message is *very* slightly different on 3.7