Qualify nested-class references when lowering pyspec class fields

Strata/Languages/Python/Specs/Decls.lean

@@ -443,6 +443,55 @@ def asTypedDict (tp : SpecType) : Option (Array DictField) := do
 
 end SpecType
 
+/-! ### Recursive identifier rewriting
+
+`mapIdentNames` applies a function `f` to every `PythonIdent` reachable
+within a `SpecType`, recursing into `SpecIdent.args` (e.g., the `Inner`
+inside `List[Inner]`) and `SpecTypedDict.fieldTypes`. The result is
+rebuilt via `union` so the "sorted by PythonIdent ordering" invariant on
+`idents` and the "sorted by field names" invariant on `typedDicts`
+survive rewrites that change an ident's ordering key (e.g., populating
+an empty `pythonModule`).
+
+Written as a mutual recursion over `SpecType`/`SpecIdent`/`SpecTypedDict`
+so each step only strips one level of nesting, keeping `termination_by
+sizeOf _` straightforward. -/
+
+mutual
+
+def SpecType.mapIdentNames (f : PythonIdent → PythonIdent) (ty : SpecType)
+    : SpecType :=
+  let identTys : Array SpecType := ty.idents.attach.map fun ⟨si, _⟩ =>
+    { SpecType.empty ty.loc with
+        idents := #[SpecIdent.mapIdentNames f si] }
+  let typedDictTys : Array SpecType := ty.typedDicts.attach.map fun ⟨td, _⟩ =>
+    { SpecType.empty ty.loc with
+        typedDicts := #[SpecTypedDict.mapIdentNames f td] }
+  let base : SpecType :=
+    { SpecType.empty ty.loc with
+        intLits := ty.intLits, stringLits := ty.stringLits }
+  (identTys ++ typedDictTys).foldl (init := base) (SpecType.union ty.loc · ·)
+termination_by sizeOf ty
+decreasing_by all_goals (cases ty; decreasing_tactic)
+
+def SpecIdent.mapIdentNames (f : PythonIdent → PythonIdent) (si : SpecIdent)
+    : SpecIdent :=
+  { name := f si.name
+    args := si.args.attach.map fun ⟨a, _⟩ => SpecType.mapIdentNames f a }
+termination_by sizeOf si
+decreasing_by cases si; decreasing_tactic
+
+def SpecTypedDict.mapIdentNames (f : PythonIdent → PythonIdent)
+    (td : SpecTypedDict) : SpecTypedDict :=
+  { fields := td.fields
+    fieldTypes := td.fieldTypes.attach.map fun ⟨a, _⟩ =>
+      SpecType.mapIdentNames f a
+    fieldRequired := td.fieldRequired }
+termination_by sizeOf td
+decreasing_by cases td; decreasing_tactic
+
+end
+
 /-- A default value for a pyspec argument.
     TODO: extend with additional constructors (e.g., string, int, bool literals)
     as PySpec gains support for richer default values. -/

Strata/Languages/Python/Specs/ToLaurel.lean

@@ -532,16 +532,54 @@ def funcDeclToLaurel (procName : String) (func : FunctionDecl)
     body := body
   }
 
-/-- Convert a class definition to Laurel types and procedures. -/
-def classDefToLaurel (cls : ClassDef) : ToLaurelM Unit := do
+/-- Rewrite each `SpecIdent` in `ty` whose `pythonModule` is empty and whose
+    `name` is a key of `subNames` to the qualified ident referring to a nested
+    class. This fixes references produced by `typeClassNoArgs` in pyspec input
+    (e.g., `inner : class(Inner)` inside `class Outer`), which lose their
+    enclosing-class context on DDM round-trip.
+
+    `subNames` maps the bare subclass name (e.g., `Inner`) to the prefixed
+    module path under which the nested class is emitted (e.g., `Outer`). Once
+    rewritten, `PythonIdent.toLaurelName` will produce `Outer_Inner`, which
+    matches the name `classDefToLaurel` assigns to the nested class. -/
+private def qualifyNestedClassRefs
+    (subNames : Std.HashMap String String) (ty : SpecType) : SpecType :=
+  if subNames.isEmpty then ty
+  else
+    ty.mapIdentNames fun id =>
+      if id.pythonModule.isEmpty then
+        match subNames[id.name]? with
+        | some qualifiedModule => { pythonModule := qualifiedModule, name := id.name }
+        | none => id
+      else id
+
+/-- Convert a class definition to Laurel types and procedures.
+
+    `isNested` is set when recursing into a parent class's `subclasses`.
+    Nested classes are addressable only via their qualified prefixed name
+    (e.g., `Outer_Inner`), so we suppress the bare-name alias registration
+    that top-level classes get — otherwise multiple sibling outers each
+    declaring a nested class of the same bare name would race to overwrite
+    the alias entry, reintroducing ambiguous resolutions. -/
+def classDefToLaurel (cls : ClassDef) (isNested : Bool := false) : ToLaurelM Unit := do
   let prefixedName ← prefixName cls.name
-  -- Register alias from unprefixed to prefixed name for type resolution
-  if prefixedName != cls.name then
+  -- Register alias from unprefixed to prefixed name for type resolution.
+  -- Only top-level classes are user-addressable by their bare name;
+  -- nested classes deliberately do not get an entry here.
+  if !isNested && prefixedName != cls.name then
     modify fun s => { s with typeAliases := s.typeAliases.insert cls.name prefixedName }
   if cls.exhaustive then
     modify fun s => { s with exhaustiveClasses := s.exhaustiveClasses.insert prefixedName }
+  -- Build a map from each nested class's bare name to the qualified module
+  -- path under which it will be emitted. Field type references that come
+  -- from `typeClassNoArgs` (and therefore lost their enclosing-class context
+  -- on DDM round-trip) are then rewritten to the qualified reference.
+  let subNames : Std.HashMap String String :=
+    cls.subclasses.foldl (init := ∅) fun m sub =>
+      m.insert sub.name prefixedName
   let laurelFields ← cls.fields.toList.mapM fun f => do
-    let ty ← specTypeToLaurelType f.type
+    let qualifiedTy := qualifyNestedClassRefs subNames f.type
+    let ty ← specTypeToLaurelType qualifiedTy
     pure { name := f.name, isMutable := true, type := ty : Laurel.Field }
   let prefixedBases := cls.bases.toList.map fun cd =>
     mkId cd.toLaurelName
@@ -554,8 +592,15 @@ def classDefToLaurel (cls : ClassDef) : ToLaurelM Unit := do
   for method in cls.methods do
     let proc ← funcDeclToLaurel (prefixedName ++ "@" ++ method.name) method (isMethod := true)
     pushProcedure proc
+  -- When recursing into nested classes, push the qualified path as the
+  -- module prefix for that subtree so the subclass's own prefixed name
+  -- aligns with how its parent references it (`subNames` above). Pass
+  -- `isNested := true` so the recursion does not register a bare alias.
   for sub in cls.subclasses do
-    classDefToLaurel sub
+    withReader
+      (fun ctx => { ctx with modulePrefix := prefixedName })
+      (classDefToLaurel sub (isNested := true))
+termination_by sizeOf cls
 decreasing_by
   · cases cls
     decreasing_tactic

StrataTest/Languages/Python/Specs/DeclsTest.lean

@@ -5,6 +5,8 @@
 -/
 import Strata.Languages.Python.Specs.Decls
 
+open Strata
+open Strata.Python
 open Strata.Python.Specs
 
 namespace DeclsTest
@@ -13,4 +15,87 @@ namespace DeclsTest
 #guard (SpecType.unionArray default
     #[SpecType.intLiteral ⟨0, 0⟩ 0, SpecType.intLiteral ⟨0, 0⟩ 0]).intLits.size == 1
 
+/-! ## `SpecType.mapIdentNames`
+
+The rewrite traverses every nested `SpecType` reachable via
+`SpecIdent.args` and `SpecTypedDict.fieldTypes`, and re-sorts the
+`idents`/`typedDicts` arrays so ordering invariants survive rewrites
+that change the ordering key. -/
+
+private def ll : SourceRange := SourceRange.none
+
+private def bareInner : PythonIdent := { pythonModule := "", name := "Inner" }
+
+/-- Rewrite bare `Inner` (empty module) to `Outer.Inner`. -/
+private def qualifyInner (id : PythonIdent) : PythonIdent :=
+  if id.pythonModule.isEmpty && id.name == "Inner" then
+    { pythonModule := "Outer", name := "Inner" }
+  else id
+
+-- Top-level bare ident is rewritten.
+/-- info: Outer.Inner -/
+#guard_msgs in
+#eval IO.println
+  (SpecType.ident ll bareInner |>.mapIdentNames qualifyInner |>.toString)
+
+-- Generic arg: `List[Inner]`. The rewrite must reach inside `SpecIdent.args`.
+/-- info: typing.List[Outer.Inner] -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.ident ll { pythonModule := "typing", name := "List" }
+      #[SpecType.ident ll bareInner]
+   |>.mapIdentNames qualifyInner).toString
+
+-- Dict args: `Dict[str, Inner]` — rewrite reaches each arg.
+/-- info: typing.Dict[builtins.str, Outer.Inner] -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.ident ll { pythonModule := "typing", name := "Dict" }
+      #[SpecType.ident ll { pythonModule := "builtins", name := "str" },
+        SpecType.ident ll bareInner]
+   |>.mapIdentNames qualifyInner).toString
+
+-- Union atom: `Union[None, Inner]`. After rewriting, `Outer.Inner` sorts
+-- ahead of `_types.NoneType` because `'O' (0x4f) < '_' (0x5f)`.
+/-- info: Union[Outer.Inner, _types.NoneType] -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.unionArray ll
+      #[SpecType.noneType ll, SpecType.ident ll bareInner]
+   |>.mapIdentNames qualifyInner).toString
+
+-- TypedDict field type: `TypedDict(f : Inner)`. The rewrite must reach
+-- `SpecTypedDict.fieldTypes`.
+/-- info: TypedDict(f : Outer.Inner) -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.typedDict ll #["f"] #[SpecType.ident ll bareInner] #[true]
+   |>.mapIdentNames qualifyInner).toString
+
+-- Nested generic: `List[TypedDict(f : Inner)]` — exercises both paths.
+/-- info: typing.List[TypedDict(f : Outer.Inner)] -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.ident ll { pythonModule := "typing", name := "List" }
+      #[SpecType.typedDict ll #["f"] #[SpecType.ident ll bareInner] #[true]]
+   |>.mapIdentNames qualifyInner).toString
+
+-- Sort restoration: rewriting an ident's `pythonModule` can change its
+-- position under `PythonIdent` ordering. The result must stay sorted so
+-- the "idents sorted by PythonIdent ordering" invariant (`Decls.lean`)
+-- holds. Before: `[a.Alpha, b.Beta]`. Rewrite `b.Beta → a.Beta`; a naive
+-- in-place map would leave them in original order — and since
+-- `SpecType.toString` walks `idents` as stored, the printed order
+-- exposes whether the sort was restored.
+/-- info: Union[a.Alpha, a.Beta] -/
+#guard_msgs in
+#eval IO.println <|
+  (SpecType.unionArray ll
+      #[SpecType.ident ll { pythonModule := "a", name := "Alpha" },
+        SpecType.ident ll { pythonModule := "b", name := "Beta" }]
+   |>.mapIdentNames (fun id =>
+      if id.pythonModule == "b" && id.name == "Beta" then
+        { pythonModule := "a", name := "Beta" }
+      else id)).toString
+
 end DeclsTest

StrataTest/Languages/Python/ToLaurelTest.lean

@@ -836,4 +836,122 @@ private def translateFunc (args : Array Arg := #[])
   -- return type assume
   assert! body.contains "assume Any..isfrom_str(result)"
 
+/-! ## Nested class field references
+
+Regression test: when a class field's type is a bare identifier (empty
+`pythonModule`) whose name matches a nested class of the enclosing class,
+`classDefToLaurel` rewrites the reference to the qualified name assigned to
+the nested class. This models PySpec inputs produced by `typeClassNoArgs`
+where the enclosing-class context is lost on DDM round-trip. -/
+
+-- Build a raw `.classDef` signature exposing `fields` and `subclasses`
+-- (the `classDef` helper above only sets methods).
+private def rawClassDef (name : String)
+    (fields : Array ClassField := #[])
+    (subclasses : Array ClassDef := #[])
+    (methods : Array FunctionDecl := #[]) : Signature :=
+  .classDef { loc, name, fields, subclasses, methods }
+
+-- Print composite types with their field names and resolved types, so the
+-- test can observe whether nested-class references were qualified.
+private def runTypeFieldsTest (sigs : Array Signature)
+    (modulePrefix : String := "") : IO Unit := do
+  let result := signaturesToLaurel "<test>" sigs modulePrefix
+  for err in result.errors do
+    IO.println s!"warning: {err.kind.phase}.{err.kind.category}: {err.message}"
+  for td in result.program.types do
+    match td with
+    | .Composite ct =>
+      IO.println s!"type {ct.name.text}"
+      for f in ct.fields do
+        IO.println s!"  {f.name.text} : {fmtHighType f.type.val}"
+    | _ =>
+      IO.println (fmtTypeDef td)
+
+-- A field whose type is a bare reference to the enclosing class's nested
+-- class is qualified to the nested class's prefixed name.
+/--
+info: type Outer
+  inner : UserDefined(Outer_Inner)
+type Outer_Inner
+-/
+#guard_msgs in
+#eval runTypeFieldsTest #[
+  rawClassDef "Outer"
+    (fields := #[{ name := "inner", type := pyClass "Inner" }])
+    (subclasses := #[{ loc, name := "Inner", methods := #[] }])
+]
+
+-- Two sibling classes each declaring a nested class of the same bare name
+-- must not collide: each outer class's bare reference is qualified with its
+-- own prefix.
+/--
+info: type A
+  child : UserDefined(A_Child)
+type A_Child
+type B
+  child : UserDefined(B_Child)
+type B_Child
+-/
+#guard_msgs in
+#eval runTypeFieldsTest #[
+  rawClassDef "A"
+    (fields := #[{ name := "child", type := pyClass "Child" }])
+    (subclasses := #[{ loc, name := "Child", methods := #[] }]),
+  rawClassDef "B"
+    (fields := #[{ name := "child", type := pyClass "Child" }])
+    (subclasses := #[{ loc, name := "Child", methods := #[] }])
+]
+
+-- A bare reference to a name that is *not* a nested class of the enclosing
+-- class is left unqualified (no spurious rewriting).
+/--
+info: type Outer
+  sibling : UserDefined(Unrelated)
+type Unrelated
+-/
+#guard_msgs in
+#eval runTypeFieldsTest #[
+  rawClassDef "Outer"
+    (fields := #[{ name := "sibling", type := pyClass "Unrelated" }]),
+  rawClassDef "Unrelated"
+]
+
+-- Nested classes must NOT leak into `typeAliases`: the map exposes
+-- unprefixed top-level class names to user code, and registering a nested
+-- class under its bare name would (a) make it user-addressable as a
+-- top-level type, and (b) let sibling outers with same-named nested
+-- classes overwrite each other.
+/--
+info: typeAliases keys: [Outer]
+-/
+#guard_msgs in
+#eval do
+  let result := signaturesToLaurel "<test>" #[
+    rawClassDef "Outer"
+      (fields := #[{ name := "inner", type := pyClass "Inner" }])
+      (subclasses := #[{ loc, name := "Inner", methods := #[] }])
+  ] "mod"
+  let keys := result.typeAliases.toArray.map Prod.fst |>.qsort (· < ·)
+  IO.println s!"typeAliases keys: {keys.toList}"
+
+-- Two sibling classes, each with a nested class of the same bare name,
+-- do not pollute `typeAliases` with the nested name (no last-writer-wins
+-- collision).
+/--
+info: typeAliases keys: [A, B]
+-/
+#guard_msgs in
+#eval do
+  let result := signaturesToLaurel "<test>" #[
+    rawClassDef "A"
+      (fields := #[{ name := "child", type := pyClass "Child" }])
+      (subclasses := #[{ loc, name := "Child", methods := #[] }]),
+    rawClassDef "B"
+      (fields := #[{ name := "child", type := pyClass "Child" }])
+      (subclasses := #[{ loc, name := "Child", methods := #[] }])
+  ] "mod"
+  let keys := result.typeAliases.toArray.map Prod.fst |>.qsort (· < ·)
+  IO.println s!"typeAliases keys: {keys.toList}"
+
 end Strata.Python.Specs.ToLaurel.Tests