[ty] Use CompactStr for StringLiteralType values

crates/ty_python_semantic/resources/mdtest/function/return_type.md

@@ -134,7 +134,6 @@ since these functions will never actually be called.
 
 ```py
 from typing import TYPE_CHECKING
-import typing
 
 if TYPE_CHECKING:
     def f() -> int: ...
@@ -200,9 +199,6 @@ if get_bool():
 if TYPE_CHECKING:
     if not TYPE_CHECKING:
         def n() -> str: ...
-
-if typing.TYPE_CHECKING:
-    def o() -> str: ...
 ```
 
 ## Conditional return type

crates/ty_python_semantic/resources/mdtest/known_constants.md

@@ -3,37 +3,27 @@
 ## `typing.TYPE_CHECKING`
 
 This constant is `True` when in type-checking mode, `False` otherwise. The symbol is defined to be
-`False` at runtime. In typeshed, it is annotated as `bool`.
+`False` at runtime. In typeshed, it is annotated as `bool`. This test makes sure that we infer
+`Literal[True]` for it anyways.
 
 ### Basic
 
 ```py
 from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    type_checking = True
-if not TYPE_CHECKING:
-    runtime = True
-
-# type_checking is treated as unconditionally assigned.
-reveal_type(type_checking)  # revealed: Literal[True]
-# error: [unresolved-reference]
-reveal_type(runtime)  # revealed: Unknown
-```
-
-### As module attribute
-
-```py
 import typing
 
-if typing.TYPE_CHECKING:
-    type_checking = True
-if not typing.TYPE_CHECKING:
-    runtime = True
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+reveal_type(typing.TYPE_CHECKING)  # revealed: Literal[True]
+```
 
-reveal_type(type_checking)  # revealed: Literal[True]
-# error: [unresolved-reference]
-reveal_type(runtime)  # revealed: Unknown
+### Aliased
+
+Make sure that we still infer the correct type if the constant has been given a different name:
+
+```py
+from typing import TYPE_CHECKING as TC
+
+reveal_type(TC)  # revealed: Literal[True]
 ```
 
 ### `typing_extensions` re-export
@@ -43,14 +33,7 @@ This should behave in the same way as `typing.TYPE_CHECKING`:
 ```py
 from typing_extensions import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    type_checking = True
-if not TYPE_CHECKING:
-    runtime = True
-
-reveal_type(type_checking)  # revealed: Literal[True]
-# error: [unresolved-reference]
-reveal_type(runtime)  # revealed: Unknown
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 ```
 
 ## User-defined `TYPE_CHECKING`
@@ -63,7 +46,7 @@ type checkers, e.g. mypy and pyright.
 
 ```py
 TYPE_CHECKING = False
-
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 if TYPE_CHECKING:
     type_checking = True
 if not TYPE_CHECKING:
@@ -78,11 +61,11 @@ reveal_type(runtime)  # revealed: Unknown
 ### With a type annotation
 
 We can also define `TYPE_CHECKING` with a type annotation. The type must be one to which `bool` can
-be assigned.
+be assigned. Even in this case, the type of `TYPE_CHECKING` is still inferred to be `Literal[True]`.
 
 ```py
 TYPE_CHECKING: bool = False
-
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 if TYPE_CHECKING:
     type_checking = True
 if not TYPE_CHECKING:
@@ -101,21 +84,6 @@ reveal_type(runtime)  # revealed: Unknown
 TYPE_CHECKING = False
 ```
 
-```py
-from constants import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    type_checking = True
-if not TYPE_CHECKING:
-    runtime = True
-
-reveal_type(type_checking)  # revealed: Literal[True]
-# error: [unresolved-reference]
-reveal_type(runtime)  # revealed: Unknown
-```
-
-### Importing user-defined `TYPE_CHECKING` from stub
-
 `stub.pyi`:
 
 ```pyi
@@ -125,16 +93,13 @@ TYPE_CHECKING: bool = ...
 ```
 
 ```py
+from constants import TYPE_CHECKING
+
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
+
 from stub import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    type_checking = True
-if not TYPE_CHECKING:
-    runtime = True
-
-reveal_type(type_checking)  # revealed: Literal[True]
-# error: [unresolved-reference]
-reveal_type(runtime)  # revealed: Unknown
+reveal_type(TYPE_CHECKING)  # revealed: Literal[True]
 ```
 
 ### Invalid assignment to `TYPE_CHECKING`
@@ -157,14 +122,12 @@ TYPE_CHECKING: int = 1
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = "str"
 
-# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = False
 
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: Literal[False] = False
 
-# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: Literal[True] = False
 ```
@@ -177,7 +140,6 @@ from typing import Literal
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str
 
-# error: [invalid-assignment]
 # error: [invalid-type-checking-constant]
 TYPE_CHECKING: str = False
 

crates/ty_python_semantic/resources/mdtest/type_api.md

@@ -154,6 +154,7 @@ the expression is not of statically known truthiness.
 
 ```py
 from ty_extensions import static_assert
+from typing import TYPE_CHECKING
 import sys
 
 static_assert(True)
@@ -173,6 +174,8 @@ static_assert("d" in "abc")  # error: "Static assertion error: argument evaluate
 n = None
 static_assert(n is None)
 
+static_assert(TYPE_CHECKING)
+
 static_assert(sys.version_info >= (3, 6))
 ```
 

crates/ty_python_semantic/src/place.rs

@@ -754,10 +754,14 @@ fn place_by_id<'db>(
             // a diagnostic if we see it being modified externally. In type inference, we
             // can assign a "narrow" type to it even if it is not *declared*. This means, we
             // do not have to call [`widen_type_for_undeclared_public_symbol`].
+            //
+            // `TYPE_CHECKING` is a special variable that should only be assigned `False`
+            // at runtime, but is always considered `True` in type checking.
+            // See mdtest/known_constants.md#user-defined-type_checking for details.
             let is_considered_non_modifiable = place_table(db, scope)
                 .place_expr(place_id)
                 .expr
-                .is_name_and(|name| matches!(name, "__slots__"));
+                .is_name_and(|name| matches!(name, "__slots__" | "TYPE_CHECKING"));
 
             if scope.file(db).is_stub(db) {
                 // We generally trust module-level undeclared places in stubs and do not union

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -549,20 +549,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
     }
 
     fn build_predicate(&mut self, predicate_node: &ast::Expr) -> PredicateOrLiteral<'db> {
-        // Some commonly used test expressions are eagerly evaluated as `true` or `false` here for
-        // performance reasons. This list does not need to be exhaustive. More complex expressions
-        // will still evaluate to the correct value during type-checking. (The one exception is
-        // `TYPE_CHECKING`; we need to detect it here in order to handle it correctly in
-        // conditions; in type inference it will resolve to its runtime value.)
+        // Some commonly used test expressions are eagerly evaluated as `true`
+        // or `false` here for performance reasons. This list does not need to
+        // be exhaustive. More complex expressions will still evaluate to the
+        // correct value during type-checking.
         fn resolve_to_literal(node: &ast::Expr) -> Option<bool> {
             match node {
                 ast::Expr::BooleanLiteral(ast::ExprBooleanLiteral { value, .. }) => Some(*value),
                 ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING" => Some(true),
-                ast::Expr::Attribute(ast::ExprAttribute { attr, .. })
-                    if attr == "TYPE_CHECKING" =>
-                {
-                    Some(true)
-                }
                 ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
                     value: ast::Number::Int(n),
                     ..
@@ -2759,12 +2753,14 @@ impl ExpressionsScopeMapBuilder {
 /// Returns if the expression is a `TYPE_CHECKING` expression.
 fn is_if_type_checking(expr: &ast::Expr) -> bool {
     matches!(expr, ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING")
-        || matches!(expr, ast::Expr::Attribute(ast::ExprAttribute { attr, .. }) if attr == "TYPE_CHECKING")
 }
 
 /// Returns if the expression is a `not TYPE_CHECKING` expression.
 fn is_if_not_type_checking(expr: &ast::Expr) -> bool {
     matches!(expr, ast::Expr::UnaryOp(ast::ExprUnaryOp { op, operand, .. }) if *op == ruff_python_ast::UnaryOp::Not
-        && is_if_type_checking(operand)
+        && matches!(
+            &**operand,
+            ast::Expr::Name(ast::ExprName { id, .. }) if id == "TYPE_CHECKING"
+        )
     )
 }

crates/ty_python_semantic/src/types.rs

@@ -6,6 +6,7 @@ use std::slice::Iter;
 
 use bitflags::bitflags;
 use call::{CallDunderError, CallError, CallErrorKind};
+use compact_str::{CompactString, ToCompactString};
 use context::InferContext;
 use diagnostic::{
     INVALID_CONTEXT_MANAGER, INVALID_SUPER_ARGUMENT, NOT_ITERABLE, POSSIBLY_UNBOUND_IMPLICIT_CALL,
@@ -16,6 +17,7 @@ use ruff_db::files::File;
 use ruff_python_ast::name::Name;
 use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_text_size::{Ranged, TextRange};
+use salsa::plumbing::interned::Lookup;
 use type_ordering::union_or_intersection_elements_ordering;
 
 pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
@@ -938,7 +940,7 @@ impl<'db> Type<'db> {
     }
 
     pub fn string_literal(db: &'db dyn Db, string: &str) -> Self {
-        Self::StringLiteral(StringLiteralType::new(db, string))
+        Self::StringLiteral(StringLiteralType::new(db, StringLiteralValue::new(string)))
     }
 
     pub fn bytes_literal(db: &'db dyn Db, bytes: &[u8]) -> Self {
@@ -5576,7 +5578,7 @@ impl<'db> Type<'db> {
             Type::SpecialForm(special_form) => Type::string_literal(db, special_form.repr()),
             Type::KnownInstance(known_instance) => Type::StringLiteral(StringLiteralType::new(
                 db,
-                known_instance.repr(db).to_string().into_boxed_str(),
+                StringLiteralValue::from(known_instance.repr(db).to_compact_string()),
             )),
             // TODO: handle more complex types
             _ => KnownClass::Str.to_instance(db),
@@ -5598,7 +5600,7 @@ impl<'db> Type<'db> {
             Type::SpecialForm(special_form) => Type::string_literal(db, special_form.repr()),
             Type::KnownInstance(known_instance) => Type::StringLiteral(StringLiteralType::new(
                 db,
-                known_instance.repr(db).to_string().into_boxed_str(),
+                StringLiteralValue::from(known_instance.repr(db).to_compact_string()),
             )),
             // TODO: handle more complex types
             _ => KnownClass::Str.to_instance(db),
@@ -8280,7 +8282,7 @@ impl<'db> IntersectionType<'db> {
 #[derive(PartialOrd, Ord)]
 pub struct StringLiteralType<'db> {
     #[returns(deref)]
-    value: Box<str>,
+    value: StringLiteralValue,
 }
 
 // The Salsa heap is tracked separately.
@@ -8297,7 +8299,41 @@ impl<'db> StringLiteralType<'db> {
     pub(crate) fn iter_each_char(self, db: &'db dyn Db) -> impl Iterator<Item = Self> {
         self.value(db)
             .chars()
-            .map(|c| StringLiteralType::new(db, c.to_string().into_boxed_str()))
+            .map(|c| StringLiteralType::new(db, StringLiteralValue::from_char(c)))
+    }
+}
+
+/// Newtype wrapper around `compact_str`'s `CompactString` so that `Lookup` can be implemented for it.
+#[derive(PartialEq, Eq, Debug, Clone, Hash, Ord, PartialOrd, get_size2::GetSize)]
+pub struct StringLiteralValue(CompactString);
+
+impl StringLiteralValue {
+    fn new(value: impl AsRef<str>) -> Self {
+        StringLiteralValue(CompactString::new(value.as_ref()))
+    }
+
+    fn from_char(c: char) -> Self {
+        StringLiteralValue(c.to_compact_string())
+    }
+}
+
+impl Lookup<StringLiteralValue> for &str {
+    fn into_owned(self) -> StringLiteralValue {
+        StringLiteralValue(CompactString::new(self))
+    }
+}
+
+impl std::ops::Deref for StringLiteralValue {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl From<compact_str::CompactString> for StringLiteralValue {
+    fn from(value: compact_str::CompactString) -> Self {
+        StringLiteralValue(value)
     }
 }
 

crates/ty_python_semantic/src/types/infer.rs

@@ -3900,7 +3900,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     ) {
                         report_invalid_type_checking_constant(&self.context, target.into());
                     }
-                    value_ty
+                    Type::BooleanLiteral(true)
                 } else if self.in_stub() && value.is_ellipsis_literal_expr() {
                     Type::unknown()
                 } else {
@@ -3989,6 +3989,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 // otherwise, assigning something other than `False` is an error
                 report_invalid_type_checking_constant(&self.context, target.into());
             }
+            declared_ty.inner = Type::BooleanLiteral(true);
         }
 
         // Handle various singletons.
@@ -4013,7 +4014,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
         if let Some(value) = value {
             let inferred_ty = self.infer_expression(value);
-            let inferred_ty = if self.in_stub() && value.is_ellipsis_literal_expr() {
+            let inferred_ty = if target
+                .as_name_expr()
+                .is_some_and(|name| &name.id == "TYPE_CHECKING")
+            {
+                Type::BooleanLiteral(true)
+            } else if self.in_stub() && value.is_ellipsis_literal_expr() {
                 declared_ty.inner_type()
             } else {
                 inferred_ty