added symbolic operations in type expressions

Author: KotlinIsland

crates/by_transforms/src/transforms/ast_driver.rs

@@ -39,8 +39,9 @@ use super::{
     float_const, force_unwrap, generic_call, generics, identity_swap, implicit_typing, init_method,
     intersection, just_float, kw_subscript, literal_types, main_function, modifiers,
     mutable_defaults, none_chain, not_type, optional_type, overload, postfix_await, propagate,
-    repeated_underscore, sentinel, some_ctor, super_keyword, top_star, tuple_index, type_is,
-    typed_dict_literal, typed_lambda, typeof_keyword, unpack, use_site_variance,
+    repeated_underscore, sentinel, some_ctor, super_keyword, symbolic_type_op, top_star,
+    tuple_index, type_is, typed_dict_literal, typed_lambda, typeof_keyword, unpack,
+    use_site_variance,
 };
 use crate::Config;
 use crate::type_info::TypeInfo;
@@ -84,6 +85,11 @@ pub(crate) struct PassContext {
     /// Lines to append AFTER the spliced body (e.g. modifiers' auto-
     /// generated `__all__ = [...]`). Driver emits each as its own line
     pub(crate) epilogue: Vec<String>,
+    /// Source ranges of operations that `symbolic_type_op` resolved up front
+    /// (e.g. `1 + 1` → `Literal[2]`). Type-aware passes skip these via
+    /// [`walk_type_positions_skipping`](super::type_expr_walker::walk_type_positions_skipping)
+    /// so they don't re-process an operation that no longer appears in the output
+    pub(crate) claimed_type_op_ranges: Vec<TextRange>,
 }
 
 /// A single AST-level rewrite pass.
@@ -286,6 +292,17 @@ pub(crate) fn run_against_source<'a>(
         text_edits: RefCell::new(vec![]),
     };
 
+    // resolve symbolic operations in type positions (`1 + 1` → `Literal[2]`)
+    // up front, from the original parse where `typeof` operands are still
+    // intact for ty to read. the pass replaces each operation node and must run
+    // before `typeof` lowering so a `typeof` operand is consumed here
+    let symbolic_folds =
+        symbolic_type_op::collect_symbolic_folds(parsed_handle.suite(), &semantic_model);
+    let symbolic_needs_literal_import = symbolic_folds.needs_literal_import;
+    let symbolic_needs_any_import = symbolic_folds.needs_any_import;
+    ctx.claimed_type_op_ranges = symbolic_folds.claimed_ranges();
+    let symbolic_pass = symbolic_type_op::SymbolicTypeOp::new(symbolic_folds);
+
     let tuple_index_inner = tuple_index::TupleIndex::new();
     let tuple_index_pass = VisitorPass {
         inner: &tuple_index_inner,
@@ -395,7 +412,12 @@ pub(crate) fn run_against_source<'a>(
         &decorator_keyword_pass,
         &unpack_pass,
         &typed_dict_literal_pass,
-        // AST-mutation passes second (may zero node ranges)
+        // AST-mutation passes second (may zero node ranges).
+        // symbolic_type_op replaces whole operation nodes (consuming any
+        // `typeof` operand) and reads original source ranges, so it must run
+        // first among the mutation passes — before `typeof` and before any
+        // pass that zeroes ranges
+        &symbolic_pass,
         &coalesce_pass,
         &cast_pass,
         &typeof_pass,
@@ -463,6 +485,18 @@ pub(crate) fn run_against_source<'a>(
         ctx.required_imports
             .push("from ty_extensions import TypeOf".to_owned());
     }
+    // symbolic folds that produced a `Literal[..]` need the import, unless the
+    // source already binds `Literal`
+    if symbolic_needs_literal_import && !literal_types::literal_already_imported(&semantic_model) {
+        ctx.required_imports
+            .push("from typing import Literal".to_owned());
+    }
+    // symbolic folds that produced `Any` (e.g. `dynamic + 1`) need the import,
+    // unless the source already binds `Any`
+    if symbolic_needs_any_import && !semantic_model.is_bound_globally("Any") {
+        ctx.required_imports
+            .push("from typing import Any".to_owned());
+    }
     if sentinel_inner.ever_changed() {
         ctx.required_imports
             .push("from typing_extensions import Sentinel".to_owned());

crates/by_transforms/src/transforms/callable.rs

@@ -518,7 +518,12 @@ impl<'ast> ruff_python_ast::visitor::Visitor<'ast> for ValueCallableWalker<'_, '
 impl TypeAwarePass for CallableSyntaxPass<'_> {
     fn run(&self, stmts: &[Stmt], types: &dyn TypeInfo, ctx: &mut PassContext) {
         let mut inner = CallableSyntax::new(self.source).with_types(types);
-        crate::transforms::type_expr_walker::walk_type_positions(stmts, Some(types), &mut inner);
+        crate::transforms::type_expr_walker::walk_type_positions_skipping(
+            stmts,
+            Some(types),
+            &ctx.claimed_type_op_ranges,
+            &mut inner,
+        );
         // also lower callable types appearing in value positions; duplicate
         // edits over type-position callables dedup in the splice
         {

crates/by_transforms/src/transforms/intersection.rs

@@ -19,7 +19,7 @@ use ruff_python_ast::{
 use ruff_text_size::{Ranged, TextRange};
 
 use super::ast_driver::{PassContext, TypeAwarePass, render_expr};
-use super::type_expr_walker::{Recurse, TypeExprVisitor, TypePos, walk_type_positions};
+use super::type_expr_walker::{Recurse, TypeExprVisitor, TypePos, walk_type_positions_skipping};
 use crate::type_info::TypeInfo;
 
 pub(crate) struct IntersectionType;
@@ -36,7 +36,7 @@ impl TypeAwarePass for IntersectionType {
             edits: Vec::new(),
             needs_import: false,
         };
-        walk_type_positions(stmts, Some(types), &mut state);
+        walk_type_positions_skipping(stmts, Some(types), &ctx.claimed_type_op_ranges, &mut state);
         ctx.text_edits.extend(state.edits);
         if state.needs_import {
             ctx.required_imports

crates/by_transforms/src/transforms/literal_types.rs

@@ -14,7 +14,9 @@ use ruff_python_ast::{Expr, ExprSubscript, Operator, Stmt, UnaryOp};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 use crate::transforms::ast_driver::{PassContext, TypeAwarePass};
-use crate::transforms::type_expr_walker::{Recurse, TypeExprVisitor, TypePos, walk_type_positions};
+use crate::transforms::type_expr_walker::{
+    Recurse, TypeExprVisitor, TypePos, walk_type_positions_skipping,
+};
 use crate::type_info::TypeInfo;
 
 pub(crate) struct LiteralType<'src> {
@@ -378,7 +380,7 @@ impl<'src> LiteralTypePass<'src> {
 impl TypeAwarePass for LiteralTypePass<'_> {
     fn run(&self, stmts: &[Stmt], types: &dyn TypeInfo, ctx: &mut PassContext) {
         let mut inner = LiteralType::new(self.source, types);
-        walk_type_positions(stmts, Some(types), &mut inner);
+        walk_type_positions_skipping(stmts, Some(types), &ctx.claimed_type_op_ranges, &mut inner);
         if inner.needs_literal_import && !literal_already_imported(types) {
             ctx.required_imports
                 .push("from typing import Literal".to_owned());

crates/by_transforms/src/transforms/mod.rs

@@ -39,6 +39,7 @@ pub(crate) mod sentinel;
 pub(crate) mod some_ctor;
 pub(crate) mod source_util;
 pub(crate) mod super_keyword;
+pub(crate) mod symbolic_type_op;
 pub(crate) mod top_star;
 pub(crate) mod tuple_index;
 pub(crate) mod type_expr_walker;