[ty] support subscripting typing.Literal with a type alias (#21207)

Fixes https://github.com/astral-sh/ty/issues/1368

## Summary

Add support for patterns like this, where a type alias to a literal type
(or union of literal types) is used to subscript `typing.Literal`:

```py
type MyAlias = Literal[1]
def _(x: Literal[MyAlias]): ...
```

This shows up in the ecosystem report for PEP 613 type alias support.

One interesting case is an alias to `bool` or an enum type. `bool` is an
equivalent type to `Literal[True, False]`, which is a union of literal
types. Similarly an enum type `E` is also equivalent to a union of its
member literal types. Since (for explicit type aliases) we infer the RHS
directly as a type expression, this makes it difficult for us to
distinguish between `bool` and `Literal[True, False]`, so we allow
either one to (or an alias to either one) to appear inside `Literal`,
where other type checkers allow only the latter.

I think for implicit type aliases it may be simpler to support only
types derived from actually subscripting `typing.Literal`, though, so I
didn't make a TODO-comment commitment here.

## Test Plan

Added mdtests, including TODO-filled tests for PEP 613 and implicit type
aliases.

### Conformance suite

All changes here are positive -- we now emit errors on lines that should
be errors. This is a side effect of the new implementation, not the
primary purpose of this PR, but it's still a positive change.

### Ecosystem

Eliminates one ecosystem false positive, where a PEP 695 type alias for
a union of literal types is used to subscript `typing.Literal`.

crates/ty_python_semantic/resources/mdtest/annotations/literal.md

@@ -39,6 +39,8 @@ def f():
     reveal_type(a7)  # revealed: None
     reveal_type(a8)  # revealed: Literal[1]
     reveal_type(b1)  # revealed: Literal[Color.RED]
+    # TODO should be `Literal[MissingT.MISSING]`
+    reveal_type(b2)  # revealed: @Todo(functional `Enum` syntax)
 
 # error: [invalid-type-form]
 invalid1: Literal[3 + 4]
@@ -66,6 +68,208 @@ a_list: list[int] = [1, 2, 3]
 invalid6: Literal[a_list[0]]
 ```
 
+## Parameterizing with a type alias
+
+`typing.Literal` can also be parameterized with a type alias for any literal type or union of
+literal types.
+
+### PEP 695 type alias
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from typing import Literal
+from enum import Enum
+
+import mod
+
+class E(Enum):
+    A = 1
+    B = 2
+
+type SingleInt = Literal[1]
+type SingleStr = Literal["foo"]
+type SingleBytes = Literal[b"bar"]
+type SingleBool = Literal[True]
+type SingleNone = Literal[None]
+type SingleEnum = Literal[E.A]
+type UnionLiterals = Literal[1, "foo", b"bar", True, None, E.A]
+# We support this because it is an equivalent type to the following union of literals, but maybe
+# we should not, because it doesn't use `Literal` form? Other type checkers do not.
+type AnEnum1 = E
+type AnEnum2 = Literal[E.A, E.B]
+# Similarly, we support this because it is equivalent to `Literal[True, False]`.
+type Bool1 = bool
+type Bool2 = Literal[True, False]
+
+def _(
+    single_int: Literal[SingleInt],
+    single_str: Literal[SingleStr],
+    single_bytes: Literal[SingleBytes],
+    single_bool: Literal[SingleBool],
+    single_none: Literal[SingleNone],
+    single_enum: Literal[SingleEnum],
+    union_literals: Literal[UnionLiterals],
+    an_enum1: Literal[AnEnum1],
+    an_enum2: Literal[AnEnum2],
+    bool1: Literal[Bool1],
+    bool2: Literal[Bool2],
+    multiple: Literal[SingleInt, SingleStr, SingleEnum],
+    single_int_other_module: Literal[mod.SingleInt],
+):
+    reveal_type(single_int)  # revealed: Literal[1]
+    reveal_type(single_str)  # revealed: Literal["foo"]
+    reveal_type(single_bytes)  # revealed: Literal[b"bar"]
+    reveal_type(single_bool)  # revealed: Literal[True]
+    reveal_type(single_none)  # revealed: None
+    reveal_type(single_enum)  # revealed: Literal[E.A]
+    reveal_type(union_literals)  # revealed: Literal[1, "foo", b"bar", True, E.A] | None
+    reveal_type(an_enum1)  # revealed: E
+    reveal_type(an_enum2)  # revealed: E
+    reveal_type(bool1)  # revealed: bool
+    reveal_type(bool2)  # revealed: bool
+    reveal_type(multiple)  # revealed: Literal[1, "foo", E.A]
+    reveal_type(single_int_other_module)  # revealed: Literal[2]
+```
+
+`mod.py`:
+
+```py
+from typing import Literal
+
+type SingleInt = Literal[2]
+```
+
+### PEP 613 type alias
+
+```py
+from typing import Literal, TypeAlias
+from enum import Enum
+
+class E(Enum):
+    A = 1
+    B = 2
+
+SingleInt: TypeAlias = Literal[1]
+SingleStr: TypeAlias = Literal["foo"]
+SingleBytes: TypeAlias = Literal[b"bar"]
+SingleBool: TypeAlias = Literal[True]
+SingleNone: TypeAlias = Literal[None]
+SingleEnum: TypeAlias = Literal[E.A]
+UnionLiterals: TypeAlias = Literal[1, "foo", b"bar", True, None, E.A]
+AnEnum1: TypeAlias = E
+AnEnum2: TypeAlias = Literal[E.A, E.B]
+Bool1: TypeAlias = bool
+Bool2: TypeAlias = Literal[True, False]
+
+def _(
+    single_int: Literal[SingleInt],
+    single_str: Literal[SingleStr],
+    single_bytes: Literal[SingleBytes],
+    single_bool: Literal[SingleBool],
+    single_none: Literal[SingleNone],
+    single_enum: Literal[SingleEnum],
+    union_literals: Literal[UnionLiterals],
+    # Could also not error
+    an_enum1: Literal[AnEnum1],  # error: [invalid-type-form]
+    an_enum2: Literal[AnEnum2],
+    # Could also not error
+    bool1: Literal[Bool1],  # error: [invalid-type-form]
+    bool2: Literal[Bool2],
+    multiple: Literal[SingleInt, SingleStr, SingleEnum],
+):
+    # TODO should be `Literal[1]`
+    reveal_type(single_int)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal["foo"]`
+    reveal_type(single_str)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[b"bar"]`
+    reveal_type(single_bytes)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[True]`
+    reveal_type(single_bool)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `None`
+    reveal_type(single_none)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[E.A]`
+    reveal_type(single_enum)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[1, "foo", b"bar", True, E.A] | None`
+    reveal_type(union_literals)  # revealed: @Todo(Inference of subscript on special form)
+    # Could also be `E`
+    reveal_type(an_enum1)  # revealed: Unknown
+    # TODO should be `E`
+    reveal_type(an_enum2)  # revealed: @Todo(Inference of subscript on special form)
+    # Could also be `bool`
+    reveal_type(bool1)  # revealed: Unknown
+    # TODO should be `bool`
+    reveal_type(bool2)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[1, "foo", E.A]`
+    reveal_type(multiple)  # revealed: @Todo(Inference of subscript on special form)
+```
+
+### Implicit type alias
+
+```py
+from typing import Literal
+from enum import Enum
+
+class E(Enum):
+    A = 1
+    B = 2
+
+SingleInt = Literal[1]
+SingleStr = Literal["foo"]
+SingleBytes = Literal[b"bar"]
+SingleBool = Literal[True]
+SingleNone = Literal[None]
+SingleEnum = Literal[E.A]
+UnionLiterals = Literal[1, "foo", b"bar", True, None, E.A]
+# For implicit type aliases, we may not want to support this. It's simpler not to, and no other
+# type checker does.
+AnEnum1 = E
+AnEnum2 = Literal[E.A, E.B]
+# For implicit type aliases, we may not want to support this.
+Bool1 = bool
+Bool2 = Literal[True, False]
+
+def _(
+    single_int: Literal[SingleInt],
+    single_str: Literal[SingleStr],
+    single_bytes: Literal[SingleBytes],
+    single_bool: Literal[SingleBool],
+    single_none: Literal[SingleNone],
+    single_enum: Literal[SingleEnum],
+    union_literals: Literal[UnionLiterals],
+    an_enum1: Literal[AnEnum1],  # error: [invalid-type-form]
+    an_enum2: Literal[AnEnum2],
+    bool1: Literal[Bool1],  # error: [invalid-type-form]
+    bool2: Literal[Bool2],
+    multiple: Literal[SingleInt, SingleStr, SingleEnum],
+):
+    # TODO should be `Literal[1]`
+    reveal_type(single_int)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal["foo"]`
+    reveal_type(single_str)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[b"bar"]`
+    reveal_type(single_bytes)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[True]`
+    reveal_type(single_bool)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `None`
+    reveal_type(single_none)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[E.A]`
+    reveal_type(single_enum)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[1, "foo", b"bar", True, E.A] | None`
+    reveal_type(union_literals)  # revealed: @Todo(Inference of subscript on special form)
+    reveal_type(an_enum1)  # revealed: Unknown
+    # TODO should be `E`
+    reveal_type(an_enum2)  # revealed: @Todo(Inference of subscript on special form)
+    reveal_type(bool1)  # revealed: Unknown
+    # TODO should be `bool`
+    reveal_type(bool2)  # revealed: @Todo(Inference of subscript on special form)
+    # TODO should be `Literal[1, "foo", E.A]`
+    reveal_type(multiple)  # revealed: @Todo(Inference of subscript on special form)
+```
+
 ## Shortening unions of literals
 
 When a Literal is parameterized with more than one value, it’s treated as exactly to equivalent to

crates/ty_python_semantic/resources/mdtest/assignment/annotations.md

@@ -259,7 +259,7 @@ class Color(Enum):
     RED = "red"
 
 f: dict[list[Literal[1]], list[Literal[Color.RED]]] = {[1]: [Color.RED, Color.RED]}
-reveal_type(f)  # revealed: dict[list[Literal[1]], list[Literal[Color.RED]]]
+reveal_type(f)  # revealed: dict[list[Literal[1]], list[Color]]
 
 class X[T]:
     def __init__(self, value: T): ...

crates/ty_python_semantic/src/types.rs

@@ -1153,6 +1153,23 @@ impl<'db> Type<'db> {
         matches!(self, Type::FunctionLiteral(..))
     }
 
+    /// Detects types which are valid to appear inside a `Literal[…]` type annotation.
+    pub(crate) fn is_literal_or_union_of_literals(&self, db: &'db dyn Db) -> bool {
+        match self {
+            Type::Union(union) => union
+                .elements(db)
+                .iter()
+                .all(|ty| ty.is_literal_or_union_of_literals(db)),
+            Type::StringLiteral(_)
+            | Type::BytesLiteral(_)
+            | Type::IntLiteral(_)
+            | Type::BooleanLiteral(_)
+            | Type::EnumLiteral(_) => true,
+            Type::NominalInstance(_) => self.is_none(db) || self.is_bool(db) || self.is_enum(db),
+            _ => false,
+        }
+    }
+
     pub(crate) fn is_union_of_single_valued(&self, db: &'db dyn Db) -> bool {
         self.as_union().is_some_and(|union| {
             union.elements(db).iter().all(|ty| {

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -6,7 +6,6 @@ use crate::types::diagnostic::{
     self, INVALID_TYPE_FORM, NON_SUBSCRIPTABLE, report_invalid_argument_number_to_special_form,
     report_invalid_arguments_to_annotated, report_invalid_arguments_to_callable,
 };
-use crate::types::enums::is_enum_class;
 use crate::types::signatures::Signature;
 use crate::types::string_annotation::parse_string_annotation;
 use crate::types::tuple::{TupleSpecBuilder, TupleType};
@@ -1369,7 +1368,6 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         parameters: &'param ast::Expr,
     ) -> Result<Type<'db>, Vec<&'param ast::Expr>> {
         Ok(match parameters {
-            // TODO handle type aliases
             ast::Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
                 let value_ty = self.infer_expression(value, TypeContext::default());
                 if matches!(value_ty, Type::SpecialForm(SpecialFormType::Literal)) {
@@ -1421,27 +1419,6 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             literal @ ast::Expr::NumberLiteral(number) if number.value.is_int() => {
                 self.infer_expression(literal, TypeContext::default())
             }
-            // For enum values
-            ast::Expr::Attribute(ast::ExprAttribute { value, attr, .. }) => {
-                let value_ty = self.infer_expression(value, TypeContext::default());
-
-                if is_enum_class(self.db(), value_ty) {
-                    let ty = value_ty
-                        .member(self.db(), &attr.id)
-                        .place
-                        .ignore_possibly_undefined()
-                        .unwrap_or(Type::unknown());
-                    self.store_expression_type(parameters, ty);
-                    ty
-                } else {
-                    self.store_expression_type(parameters, Type::unknown());
-                    if value_ty.is_todo() {
-                        value_ty
-                    } else {
-                        return Err(vec![parameters]);
-                    }
-                }
-            }
             // for negative and positive numbers
             ast::Expr::UnaryOp(u)
                 if matches!(u.op, ast::UnaryOp::USub | ast::UnaryOp::UAdd)
@@ -1451,6 +1428,35 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                 self.store_expression_type(parameters, ty);
                 ty
             }
+            // enum members and aliases to literal types
+            ast::Expr::Name(_) | ast::Expr::Attribute(_) => {
+                let subscript_ty = self.infer_expression(parameters, TypeContext::default());
+                // TODO handle implicit type aliases also
+                match subscript_ty {
+                    // type aliases to literal types
+                    Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => {
+                        let value_ty = type_alias.value_type(self.db());
+                        if value_ty.is_literal_or_union_of_literals(self.db()) {
+                            return Ok(value_ty);
+                        }
+                    }
+                    // `Literal[SomeEnum.Member]`
+                    Type::EnumLiteral(_) => {
+                        return Ok(subscript_ty);
+                    }
+                    // `Literal[SingletonEnum.Member]`, where `SingletonEnum.Member` simplifies to
+                    // just `SingletonEnum`.
+                    Type::NominalInstance(_) if subscript_ty.is_enum(self.db()) => {
+                        return Ok(subscript_ty);
+                    }
+                    // suppress false positives for e.g. members of functional-syntax enums
+                    Type::Dynamic(DynamicType::Todo(_)) => {
+                        return Ok(subscript_ty);
+                    }
+                    _ => {}
+                }
+                return Err(vec![parameters]);
+            }
             _ => {
                 self.infer_expression(parameters, TypeContext::default());
                 return Err(vec![parameters]);