use visitor.visit(...) in Type methods instead of NewType methods

Author: oconnor663

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

@@ -130,9 +130,7 @@ reveal_type(d)  # revealed: D
 ```
 
 Normal classes can't inherit from newtypes, but generic classes can be parametrized with them, so we
-also need to detect "ordinary" type cycles that happen to involve a newtype. (This turns out to be
-tricky in the implementation, see comments around `NewType::normalized_impl` and
-`NewType::apply_type_mapping_impl`.)
+also need to detect "ordinary" type cycles that happen to involve a newtype.
 
 ```py
 E = NewType("E", list["E"])

crates/ty_python_semantic/src/types.rs

@@ -1379,7 +1379,11 @@ impl<'db> Type<'db> {
             }
             Type::TypeAlias(alias) => alias.value_type(db).normalized_impl(db, visitor),
             Type::NewTypeInstance(newtype) => {
-                Type::NewTypeInstance(newtype.normalized_impl(db, visitor))
+                visitor.visit(self, || {
+                    Type::NewTypeInstance(newtype.map_base_class_type(db, |class_type| {
+                        class_type.normalized_impl(db, visitor)
+                    }))
+                })
             }
             Type::LiteralString
             | Type::AlwaysFalsy
@@ -6892,9 +6896,11 @@ impl<'db> Type<'db> {
                 instance.apply_type_mapping_impl(db, type_mapping, tcx, visitor)
             },
 
-            Type::NewTypeInstance(newtype) => {
-                Type::NewTypeInstance(newtype.apply_type_mapping_impl(db, type_mapping, tcx, visitor))
-            },
+            Type::NewTypeInstance(newtype) => visitor.visit(self, || {
+                Type::NewTypeInstance(newtype.map_base_class_type(db, |class_type| {
+                    class_type.apply_type_mapping_impl(db, type_mapping, tcx, visitor)
+                }))
+            }),
 
             Type::ProtocolInstance(instance) => {
                 // TODO: Add tests for materialization once subtyping/assignability is implemented for
@@ -7754,7 +7760,10 @@ impl<'db> KnownInstanceType<'db> {
                 // Nothing to normalize
                 Self::ConstraintSet(set)
             }
-            Self::NewType(newtype) => Self::NewType(newtype.normalized_impl(db, visitor)),
+            Self::NewType(newtype) => Self::NewType(
+                newtype
+                    .map_base_class_type(db, |class_type| class_type.normalized_impl(db, visitor)),
+            ),
         }
     }
 

crates/ty_python_semantic/src/types/newtype.rs

@@ -3,10 +3,7 @@ use std::collections::BTreeSet;
 use crate::Db;
 use crate::semantic_index::definition::{Definition, DefinitionKind};
 use crate::types::constraints::ConstraintSet;
-use crate::types::{
-    ApplyTypeMappingVisitor, ClassType, NormalizedVisitor, Type, TypeContext, TypeMapping,
-    definition_expression_type,
-};
+use crate::types::{ClassType, Type, definition_expression_type};
 use ruff_db::parsed::parsed_module;
 use ruff_python_ast as ast;
 
@@ -133,7 +130,10 @@ impl<'db> NewType<'db> {
         ConstraintSet::from(!self.is_subtype_of(db, other) && !other.is_subtype_of(db, self))
     }
 
-    fn map_base_class_type(
+    /// Create a new `NewType` by mapping the underlying `ClassType`. This descends through any
+    /// number of nested `NewType` layers and rebuilds the whole chain. In the rare case of cyclic
+    /// `NewType`s with no underlying `ClassType`, this has no effect and does not call `f`.
+    pub(crate) fn map_base_class_type(
         self,
         db: &'db dyn Db,
         f: impl FnOnce(ClassType<'db>) -> ClassType<'db>,
@@ -149,10 +149,11 @@ impl<'db> NewType<'db> {
         for base in self.iter_bases(db) {
             match base {
                 // Build up the stack of intermediate newtypes that we'll need to re-wrap after
-                // we've mapped the base class.
+                // we've mapped the `ClassType`.
                 NewTypeBase::NewType(base_newtype) => inner_newtype_stack.push(base_newtype),
+                // We've reached the `ClassType`.
                 NewTypeBase::ClassType(base_class_type) => {
-                    // We've reached the base class. Call `f`.
+                    // Call `f`.
                     let mut mapped_base = NewTypeBase::ClassType(f(base_class_type));
                     // Re-wrap the mapped base class in however many newtypes we unwrapped.
                     for inner_newtype in inner_newtype_stack.into_iter().rev() {
@@ -172,74 +173,10 @@ impl<'db> NewType<'db> {
                 }
             }
         }
-        // If we get here, there is no base class (because this newtype is cyclic), and we don't
+        // If we get here, there is no `ClassType` (because this newtype is cyclic), and we don't
         // call `f` at all.
         self
     }
-
-    pub(crate) fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
-        self.map_base_class_type(db, |base_class_type| {
-            // We don't really want to work with `Type` directly (`NewTypeBase` is much more
-            // restricted), but we need to let the `NormalizedVisitor` see the `Type` of the base
-            // class, because this is its only chance to detect an indirect cycle like this one:
-            //
-            // ```py
-            // Foo = NewType("Foo", list["Foo"])
-            // ```
-            //
-            // Note that NewTypeBaseIter does *not* catch this kind of cycle, because it stops
-            // after returning `NewTypeBase::ClassType` for `list`.
-            let mut is_first_encounter = false;
-            let base_type = Type::instance(db, base_class_type);
-            _ = visitor.visit(base_type, || {
-                // If we hit in a cycle, this closure won't run.
-                is_first_encounter = true;
-                base_type
-            });
-            if is_first_encounter {
-                base_class_type.normalized_impl(db, visitor)
-            } else {
-                // Cycle detected. Fall back to `object`.
-                // REVIEWERS: Is this correct?
-                ClassType::object(db)
-            }
-        })
-    }
-
-    pub(crate) fn apply_type_mapping_impl<'a>(
-        self,
-        db: &'db dyn Db,
-        type_mapping: &TypeMapping<'a, 'db>,
-        tcx: TypeContext<'db>,
-        visitor: &ApplyTypeMappingVisitor<'db>,
-    ) -> Self {
-        self.map_base_class_type(db, |base_class_type| {
-            // We don't really want to work with `Type` directly (`NewTypeBase` is much more
-            // restricted), but we need to let the `ApplyTypeMappingVisitor` see the `Type` of the
-            // base class, because this is its only chance to detect an indirect cycle like:
-            //
-            // ```py
-            // Foo = NewType("Foo", list["Foo"])
-            // ```
-            //
-            // Note that NewTypeBaseIter does *not* catch this kind of cycle, because it stops
-            // after returning `NewTypeBase::ClassType` for `list`.
-            let mut is_first_encounter = false;
-            let base_type = Type::instance(db, base_class_type);
-            _ = visitor.visit(base_type, || {
-                // If we hit in a cycle, this closure won't run.
-                is_first_encounter = true;
-                base_type
-            });
-            if is_first_encounter {
-                base_class_type.apply_type_mapping_impl(db, type_mapping, tcx, visitor)
-            } else {
-                // Cycle detected. Fall back to `object`.
-                // REVIEWERS: Is this correct?
-                ClassType::object(db)
-            }
-        })
-    }
 }
 
 /// `typing.NewType` typically wraps a class type, but it can also wrap another newtype.