Add gaps identified by the reviewer and add tests to prevent regression

Alan4506 · Alan4506 · commit d0f82cc65051 · 2026-05-13T02:12:01.000-04:00
diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/writer/ImportDeclarations.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/writer/ImportDeclarations.java
@@ -108,16 +108,34 @@ private ImportDeclarations addImportToMap(
      * Retroactively aliases an existing import to avoid name collisions.
      *
      * <p>This is called by {@link PythonWriter} at {@code toString()} time when a collision
-     * is detected between an imported name and a locally-defined name. The import statement
-     * is rewritten from {@code from namespace import name} to {@code from namespace import name as alias}.
+     * is detected. The import statement is rewritten from {@code from namespace import name}
+     * to {@code from namespace import name as alias}.
      *
-     * @param namespace The module namespace of the import.
+     * <p>The lookup mirrors the dispatch logic in {@link #addImport} so that aliasing works
+     * for all three import maps:
+     * <ul>
+     *   <li>{@code stdlibImports} — keyed by absolute namespace.</li>
+     *   <li>{@code externalImports} — keyed by absolute namespace.</li>
+     *   <li>{@code localImports} — keyed by relativized namespace for same-package imports.</li>
+     * </ul>
+     *
+     * @param namespace The module namespace of the import (absolute).
      * @param name The original imported name.
      * @param alias The alias to use instead.
      */
     void aliasImport(String namespace, String name, String alias) {
+        // stdlib imports are stored with absolute namespaces.
+        aliasImportInMap(namespace, name, alias, stdlibImports);
+
+        // external imports are stored with absolute namespaces.
         aliasImportInMap(namespace, name, alias, externalImports);
-        aliasImportInMap(namespace, name, alias, localImports);
+
+        // local imports may be stored with relativized namespaces (mirroring addImport).
+        if (namespace.startsWith(settings.moduleName())) {
+            var isTestModule = this.localNamespace.startsWith("tests");
+            var ns = isTestModule ? namespace : relativize(namespace);
+            aliasImportInMap(ns, name, alias, localImports);
+        }
     }
 
     private void aliasImportInMap(
diff --git a/codegen/core/src/main/java/software/amazon/smithy/python/codegen/writer/PythonWriter.java b/codegen/core/src/main/java/software/amazon/smithy/python/codegen/writer/PythonWriter.java
@@ -479,15 +479,29 @@ public String apply(Object type, String indent) {
         /**
          * Returns a placeholder token for the symbol and registers it in the symbol table.
          *
-         * <p>Only symbols that are imported from another module (have a namespace but no
-         * definition file) get a placeholder. Symbols without a namespace (builtins like
-         * int, str, bool) and symbols with a definition file (locally-defined types) are
-         * returned as-is since they don't need collision detection.
+         * <p>Symbols that do not participate in collision detection are returned as-is:
+         * <ul>
+         *   <li>Builtins (no namespace) such as {@code int}, {@code str}, {@code bool}:
+         *       they produce no import statement, so there is no import we could alias
+         *       to resolve a collision. A sanity check across all public AWS service
+         *       models confirms no class-generating shape name collides with a Python
+         *       built-in emitted by the codegen, so builtins are not special-cased here.</li>
+         *   <li>Symbols defined in the current writer's file (namespace equals the
+         *       writer's package name): they produce no import and cannot collide
+         *       with themselves.</li>
+         * </ul>
+         *
+         * <p>All other symbols (framework types and generated symbols imported from
+         * other files in the same package) are registered in the symbol table so
+         * collisions can be detected at {@link PythonWriter#toString()} time.
          */
         private String resolvePlaceholder(Symbol symbol) {
-            if (symbol.getNamespace().isEmpty() || !symbol.getDefinitionFile().isEmpty()) {
-                // No namespace means builtin. Has definition file means locally-defined.
-                // Neither case needs collision detection.
+            if (symbol.getNamespace().isEmpty()) {
+                // Builtin — no import statement is produced, so there is no alias to apply.
+                return symbol.getName();
+            }
+            if (symbol.getNamespace().equals(fullPackageName)) {
+                // Defined in the current writer's file — no import needed.
                 return symbol.getName();
             }
 
diff --git a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/PythonSymbolProviderTest.java
@@ -0,0 +1,99 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.UnionShape;
+
+public class PythonSymbolProviderTest {
+
+    private static final String NS = "smithy.example";
+
+    @Test
+    public void testUnionMemberVariantNameCollidingWithShapeUsesUnderscoreSeparator() {
+        Model model = loadModel("""
+                $version: "2"
+                namespace smithy.example
+
+                service TestService {
+                    version: "2024-01-01"
+                    operations: [TestOp]
+                }
+
+                operation TestOp {
+                    input: TestOpInput
+                }
+
+                structure TestOpInput {
+                    principal: Principal
+                }
+
+                union Principal {
+                    user: PrincipalUser
+                }
+
+                structure PrincipalUser {
+                    name: String
+                }
+                """);
+        PythonSymbolProvider provider = createProvider(model);
+        var userMember = model.expectShape(ShapeId.from(NS + "#Principal$user"), MemberShape.class);
+
+        assertEquals("Principal_User", provider.toSymbol(userMember).getName());
+    }
+
+    @Test
+    public void testUnionUnknownVariantNameCollidingWithShapeUsesUnderscoreSeparator() {
+        Model model = loadModel("""
+                $version: "2"
+                namespace smithy.example
+
+                service TestService {
+                    version: "2024-01-01"
+                    operations: [TestOp]
+                }
+
+                operation TestOp {
+                    input: TestOpInput
+                }
+
+                structure TestOpInput {
+                    value: MyUnion
+                    other: MyUnionUnknown
+                }
+
+                union MyUnion {
+                    foo: String
+                }
+
+                structure MyUnionUnknown {
+                    message: String
+                }
+                """);
+        PythonSymbolProvider provider = createProvider(model);
+        var union = model.expectShape(ShapeId.from(NS + "#MyUnion"), UnionShape.class);
+
+        assertEquals("MyUnion_Unknown",
+                provider.toSymbol(union).expectProperty(SymbolProperties.UNION_UNKNOWN).getName());
+    }
+
+    private static Model loadModel(String smithyIdl) {
+        return Model.assembler().addUnparsedModel("test.smithy", smithyIdl).assemble().unwrap();
+    }
+
+    private static PythonSymbolProvider createProvider(Model model) {
+        PythonSettings settings = PythonSettings.builder()
+                .service(ShapeId.from(NS + "#TestService"))
+                .moduleName("test_client")
+                .moduleVersion("0.0.1")
+                .build();
+        return new PythonSymbolProvider(model, settings);
+    }
+}
diff --git a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/writer/ImportDeclarationsTest.java b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/writer/ImportDeclarationsTest.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.writer;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.python.codegen.PythonSettings;
+
+public class ImportDeclarationsTest {
+
+    private static final String MODULE = "aws_sdk_example";
+    private static final String LOCAL_NAMESPACE = MODULE + ".config";
+
+    @Test
+    public void testAliasImportRewritesStdlibImport() {
+        ImportDeclarations imports = createImports();
+        imports.addStdlibImport("decimal", "Decimal");
+
+        imports.aliasImport("decimal", "Decimal", "_Decimal");
+
+        String out = imports.toString();
+        assertTrue(normalize(out).contains("from decimal import Decimal as _Decimal"));
+    }
+
+    @Test
+    public void testAliasImportRewritesExternalImport() {
+        ImportDeclarations imports = createImports();
+        imports.addImport("smithy_core.documents", "Document", "Document");
+
+        imports.aliasImport("smithy_core.documents", "Document", "_Document");
+
+        String out = imports.toString();
+        assertTrue(normalize(out).contains("from smithy_core.documents import Document as _Document"));
+    }
+
+    @Test
+    public void testAliasImportRewritesLocalImport() {
+        ImportDeclarations imports = createImports();
+        imports.addImport(MODULE + ".models", "Foo", "Foo");
+
+        imports.aliasImport(MODULE + ".models", "Foo", "_Foo");
+
+        String out = imports.toString();
+        assertTrue(normalize(out).contains("from .models import Foo as _Foo"));
+    }
+
+    @Test
+    public void testAliasImportDoesNotOverwritePreExistingAlias() {
+        ImportDeclarations imports = createImports();
+        imports.addImport("smithy_core.documents", "Document", "OriginalAlias");
+
+        imports.aliasImport("smithy_core.documents", "Document", "_ShouldNotApply");
+
+        String out = imports.toString();
+        String normalized = normalize(out);
+        assertTrue(normalized.contains("from smithy_core.documents import Document as OriginalAlias"));
+        assertFalse(normalized.contains("_ShouldNotApply"));
+    }
+
+    private static ImportDeclarations createImports() {
+        PythonSettings settings = mock(PythonSettings.class);
+        when(settings.moduleName()).thenReturn(MODULE);
+        return new ImportDeclarations(settings, LOCAL_NAMESPACE);
+    }
+
+    private static String normalize(String output) {
+        return output.replaceAll("\\(\\s+", "")
+                .replaceAll(",\\s*\\)", "")
+                .replaceAll("\\s+", " ");
+    }
+}
diff --git a/codegen/core/src/test/java/software/amazon/smithy/python/codegen/writer/PythonWriterTest.java b/codegen/core/src/test/java/software/amazon/smithy/python/codegen/writer/PythonWriterTest.java
@@ -0,0 +1,109 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.codegen.writer;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.python.codegen.PythonSettings;
+
+public class PythonWriterTest {
+
+    private static final String CURRENT_PACKAGE = "aws_sdk_example.models";
+
+    @Test
+    public void testFrameworkSymbolCollidingWithLocalIsAliased() {
+        PythonWriter writer = createWriter(CURRENT_PACKAGE);
+        Symbol framework = frameworkSymbol("smithy_core.documents", "Document");
+        Symbol local = generatedSymbol(CURRENT_PACKAGE, "Document");
+        writer.addLocallyDefinedSymbol(local);
+
+        writer.write("value: $T", framework);
+        String out = writer.toString();
+
+        assertTrue(normalize(out).contains("from smithy_core.documents import Document as _Document"));
+        assertTrue(out.contains("value: _Document"));
+    }
+
+    @Test
+    public void testTwoFrameworkSymbolsWithSameSimpleNameGetModuleAliases() {
+        PythonWriter writer = createWriter("aws_sdk_example.auth");
+        Symbol a = frameworkSymbol("smithy_core.auth", "AuthOption");
+        Symbol b = frameworkSymbol("smithy_core.interfaces.auth", "AuthOption");
+
+        writer.write("x: $T", a);
+        writer.write("y: $T", b);
+        String out = writer.toString();
+        String normalized = normalize(out);
+
+        assertTrue(normalized.contains(
+                "from smithy_core.auth import AuthOption as _smithy_core_auth_AuthOption"));
+        assertTrue(normalized.contains(
+                "from smithy_core.interfaces.auth import "
+                        + "AuthOption as _smithy_core_interfaces_auth_AuthOption"));
+        assertTrue(out.contains("x: _smithy_core_auth_AuthOption"));
+        assertTrue(out.contains("y: _smithy_core_interfaces_auth_AuthOption"));
+    }
+
+    @Test
+    public void testCrossFileGeneratedSymbolCollidingWithFrameworkImportIsAliased() {
+        PythonWriter writer = createWriter("aws_sdk_example.config");
+        Symbol framework = frameworkSymbol("smithy_core.aio.interfaces", "EndpointResolver");
+        Symbol crossFileGenerated = generatedSymbol("aws_sdk_example.models", "EndpointResolver");
+
+        writer.write("x: $T", framework);
+        writer.write("y: $T", crossFileGenerated);
+        String out = writer.toString();
+        String normalized = normalize(out);
+
+        assertTrue(normalized.contains(
+                "from smithy_core.aio.interfaces import "
+                        + "EndpointResolver as _smithy_core_aio_interfaces_EndpointResolver"));
+        assertTrue(normalized.contains(
+                "from .models import EndpointResolver as _aws_sdk_example_models_EndpointResolver"));
+    }
+
+    @Test
+    public void testSymbolDefinedInCurrentWriterFileIsNotRegistered() {
+        PythonWriter writer = createWriter(CURRENT_PACKAGE);
+        Symbol selfReference = generatedSymbol(CURRENT_PACKAGE, "MyStruct");
+        writer.addLocallyDefinedSymbol(selfReference);
+
+        writer.write("value: $T", selfReference);
+        String out = writer.toString();
+
+        assertFalse(out.contains("import MyStruct"));
+        assertFalse(out.contains("_MyStruct"));
+        assertTrue(out.contains("value: MyStruct"));
+    }
+
+    private static PythonWriter createWriter(String fullPackageName) {
+        PythonSettings settings = mock(PythonSettings.class);
+        when(settings.moduleName()).thenReturn(fullPackageName.split("\\.")[0]);
+        return new PythonWriter(settings, fullPackageName);
+    }
+
+    private static Symbol frameworkSymbol(String namespace, String name) {
+        return Symbol.builder().name(name).namespace(namespace, ".").build();
+    }
+
+    private static Symbol generatedSymbol(String namespace, String name) {
+        return Symbol.builder()
+                .name(name)
+                .namespace(namespace, ".")
+                .definitionFile("./src/" + namespace.replace('.', '/') + ".py")
+                .build();
+    }
+
+    private static String normalize(String output) {
+        return output.replaceAll("\\(\\s+", "")
+                .replaceAll(",\\s*\\)", "")
+                .replaceAll("\\s+", " ");
+    }
+}
