Refactor and move config resolution codegen logic from generic layer to AWS layer

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsConfigPropertyMetadata.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Optional;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * AWS-specific config resolution metadata for a config property.
+ * Holds validators, custom resolvers, and default values.
+ */
+@SmithyInternalApi
+public record AwsConfigPropertyMetadata(
+        Symbol validator,
+        Symbol customResolver,
+        String defaultValue
+) {
+    public Optional<Symbol> validatorOpt() {
+        return Optional.ofNullable(validator);
+    }
+
+    public Optional<Symbol> customResolverOpt() {
+        return Optional.ofNullable(customResolver);
+    }
+
+    public Optional<String> defaultValueOpt() {
+        return Optional.ofNullable(defaultValue);
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public static final class Builder {
+        private Symbol validator;
+        private Symbol customResolver;
+        private String defaultValue;
+
+        public Builder validator(Symbol validator) {
+            this.validator = validator;
+            return this;
+        }
+
+        public Builder customResolver(Symbol customResolver) {
+            this.customResolver = customResolver;
+            return this;
+        }
+
+        public Builder defaultValue(String defaultValue) {
+            this.defaultValue = defaultValue;
+            return this;
+        }
+
+        public AwsConfigPropertyMetadata build() {
+            return new AwsConfigPropertyMetadata(validator, customResolver, defaultValue);
+        }
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsConfigResolutionIntegration.java

@@ -0,0 +1,254 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import software.amazon.smithy.python.codegen.ConfigProperty;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.sections.ConfigSection;
+import software.amazon.smithy.python.codegen.writer.PythonWriter;
+import software.amazon.smithy.utils.CodeInterceptor;
+import software.amazon.smithy.utils.CodeSection;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+/**
+ * Intercepts the generated Config class to add AWS-specific descriptor-based
+ * config resolution, keeping the generic ConfigGenerator unchanged.
+ */
+@SmithyInternalApi
+public class AwsConfigResolutionIntegration implements PythonIntegration {
+
+    // Metadata for properties that use descriptors, keyed by property name.
+    private static final Map<String, AwsConfigPropertyMetadata> DESCRIPTOR_METADATA = Map.of(
+            "region", AwsConfiguration.REGION_METADATA,
+            "retry_strategy", AwsConfiguration.RETRY_STRATEGY_METADATA,
+            "sdk_ua_app_id", AwsUserAgentIntegration.SDK_UA_APP_ID_METADATA
+    );
+
+    @Override
+    public List<? extends CodeInterceptor<? extends CodeSection, PythonWriter>> interceptors(
+            GenerationContext context
+    ) {
+        return List.of(
+                new PropertyDeclarationInterceptor(context),
+                new PropertyInitInterceptor(),
+                new PreDeclarationsInterceptor(),
+                new PreInitInterceptor(),
+                new ConfigTailInterceptor()
+        );
+    }
+
+    // Replaces plain field declarations with descriptor assignments for properties
+    // that have AWS metadata registered.
+    private static final class PropertyDeclarationInterceptor
+            implements CodeInterceptor<ConfigSection.PropertyDeclarationSection, PythonWriter> {
+
+        private final GenerationContext context;
+
+        PropertyDeclarationInterceptor(GenerationContext context) {
+            this.context = context;
+        }
+
+        @Override
+        public Class<ConfigSection.PropertyDeclarationSection> sectionType() {
+            return ConfigSection.PropertyDeclarationSection.class;
+        }
+
+        @Override
+        public void write(
+                PythonWriter writer,
+                String previousText,
+                ConfigSection.PropertyDeclarationSection section
+        ) {
+            ConfigProperty prop = section.property();
+            AwsConfigPropertyMetadata meta = DESCRIPTOR_METADATA.get(prop.name());
+
+            if (meta == null) {
+                writer.write(previousText.stripTrailing());
+                return;
+            }
+
+            String typeHint = prop.type().getName();
+            if (prop.isNullable() && !typeHint.endsWith("| None")) {
+                typeHint = typeHint + " | None";
+            }
+            writer.write("$L: $L = _descriptors['$L']  # type: ignore[assignment]",
+                    prop.name(), typeHint, prop.name());
+            writer.writeDocs(prop.documentation(), context);
+        }
+    }
+
+    // Skips `self.X = X` initialization for descriptor properties since the
+    // descriptor handles resolution.
+    private static final class PropertyInitInterceptor
+            implements CodeInterceptor<ConfigProperty.InitializeConfigPropertySection, PythonWriter> {
+
+        @Override
+        public Class<ConfigProperty.InitializeConfigPropertySection> sectionType() {
+            return ConfigProperty.InitializeConfigPropertySection.class;
+        }
+
+        @Override
+        public void write(
+                PythonWriter writer,
+                String previousText,
+                ConfigProperty.InitializeConfigPropertySection section
+        ) {
+            if (DESCRIPTOR_METADATA.containsKey(section.property().name())) {
+                return;
+            }
+            writer.write(previousText.stripTrailing());
+        }
+    }
+
+    // Injects _descriptors dict and _resolver field before property declarations.
+    private static final class PreDeclarationsInterceptor
+            implements CodeInterceptor<ConfigSection.PrePropertyDeclarationsSection, PythonWriter> {
+
+        @Override
+        public Class<ConfigSection.PrePropertyDeclarationsSection> sectionType() {
+            return ConfigSection.PrePropertyDeclarationsSection.class;
+        }
+
+        @Override
+        public void write(
+                PythonWriter writer,
+                String previousText,
+                ConfigSection.PrePropertyDeclarationsSection section
+        ) {
+            List<ConfigProperty> descriptorProps = section.properties().stream()
+                    .filter(p -> DESCRIPTOR_METADATA.containsKey(p.name()))
+                    .collect(Collectors.toList());
+
+            if (descriptorProps.isEmpty()) {
+                return;
+            }
+
+            addImports(writer, descriptorProps);
+
+            writer.write("# Config properties using descriptors");
+            writer.openBlock("_descriptors = {");
+            for (ConfigProperty prop : descriptorProps) {
+                AwsConfigPropertyMetadata meta = DESCRIPTOR_METADATA.get(prop.name());
+                StringBuilder sb = new StringBuilder();
+                sb.append("'").append(prop.name()).append("': ConfigProperty('")
+                        .append(prop.name()).append("'");
+                if (meta != null) {
+                    meta.validatorOpt().ifPresent(sym ->
+                            sb.append(", validator=").append(sym.getName()));
+                    meta.customResolverOpt().ifPresent(sym ->
+                            sb.append(", resolver_func=").append(sym.getName()));
+                    meta.defaultValueOpt().ifPresent(val ->
+                            sb.append(", default_value=").append(val));
+                }
+                sb.append("),");
+                writer.write(sb.toString());
+            }
+            writer.closeBlock("}");
+            writer.write("");
+            writer.write("_resolver: ConfigResolver");
+        }
+
+        private void addImports(PythonWriter writer, List<ConfigProperty> descriptorProps) {
+            writer.addImport("smithy_aws_core.config.property", "ConfigProperty");
+            writer.addImport("smithy_aws_core.config.resolver", "ConfigResolver");
+            writer.addImport("smithy_aws_core.config.sources", "EnvironmentSource");
+            writer.addImport("smithy_aws_core.config.source_info", "SourceInfo");
+
+            for (ConfigProperty prop : descriptorProps) {
+                AwsConfigPropertyMetadata meta = DESCRIPTOR_METADATA.get(prop.name());
+                if (meta == null) {
+                    continue;
+                }
+                meta.validatorOpt().ifPresent(sym ->
+                        writer.addImport(sym.getNamespace(), sym.getName()));
+                meta.customResolverOpt().ifPresent(sym ->
+                        writer.addImport(sym.getNamespace(), sym.getName()));
+                meta.defaultValueOpt().ifPresent(val -> {
+                    if (val.contains("RetryStrategyOptions")) {
+                        writer.addImport("smithy_core.retries", "RetryStrategyOptions");
+                    }
+                });
+            }
+        }
+    }
+
+    // Injects resolver initialization at the start of __init__.
+    private static final class PreInitInterceptor
+            implements CodeInterceptor<ConfigSection.PreInitializePropertiesSection, PythonWriter> {
+
+        @Override
+        public Class<ConfigSection.PreInitializePropertiesSection> sectionType() {
+            return ConfigSection.PreInitializePropertiesSection.class;
+        }
+
+        @Override
+        public void write(
+                PythonWriter writer,
+                String previousText,
+                ConfigSection.PreInitializePropertiesSection section
+        ) {
+            boolean hasDescriptors = section.properties().stream()
+                    .anyMatch(p -> DESCRIPTOR_METADATA.containsKey(p.name()));
+
+            if (!hasDescriptors) {
+                return;
+            }
+
+            writer.write("self._resolver = ConfigResolver(sources=[EnvironmentSource()])");
+            writer.write("");
+            writer.write("# Only set if provided (not None) to allow resolution from sources");
+            writer.write("for key in self.__class__._descriptors.keys():");
+            writer.indent();
+            writer.write("value = locals().get(key)");
+            writer.write("if value is not None:");
+            writer.indent();
+            writer.write("setattr(self, key, value)");
+            writer.dedent();
+            writer.dedent();
+        }
+    }
+
+    // Appends get_source() method to the Config class.
+    private static final class ConfigTailInterceptor
+            implements CodeInterceptor<ConfigSection, PythonWriter> {
+
+        @Override
+        public Class<ConfigSection> sectionType() {
+            return ConfigSection.class;
+        }
+
+        @Override
+        public void write(PythonWriter writer, String previousText, ConfigSection section) {
+            boolean hasDescriptors = section.properties().stream()
+                    .anyMatch(p -> DESCRIPTOR_METADATA.containsKey(p.name()));
+
+            writer.write(previousText.stripTrailing());
+
+            if (!hasDescriptors) {
+                return;
+            }
+
+            writer.write("""
+
+                    def get_source(self, key: str) -> SourceInfo | None:
+                        \"""Get the source that provided a configuration value.
+
+                        Args:
+                            key: The configuration key (e.g., 'region', 'retry_strategy')
+
+                        Returns:
+                            The source info (SimpleSource or ComplexSource),
+                            or None if the key hasn't been resolved yet.
+                        \"""
+                        cached = self.__dict__.get(f'_cache_{key}')
+                        return cached[1] if cached else None
+                """);
+        }
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsConfiguration.java

@@ -17,22 +17,16 @@ private AwsConfiguration() {}
 
     public static final ConfigProperty REGION = ConfigProperty.builder()
             .name("region")
-            .type(Symbol.builder().name("str | None").build())
-            .documentation(" The AWS region to connect to. The configured region is used to "
+            .type(Symbol.builder().name("str").build())
+            .documentation("The AWS region to connect to. The configured region is used to "
                     + "determine the service endpoint.")
             .nullable(false)
-            .useDescriptor(true)
-            .validator(Symbol.builder()
-                    .name("validate_region")
-                    .namespace("smithy_aws_core.config.validators", ".")
-                    .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
-                    .build())
             .build();
 
     public static final ConfigProperty RETRY_STRATEGY = ConfigProperty.builder()
             .name("retry_strategy")
             .type(Symbol.builder()
-                    .name("RetryStrategy | RetryStrategyOptions | None")
+                    .name("RetryStrategy | RetryStrategyOptions")
                     .addReference(Symbol.builder()
                             .name("RetryStrategy")
                             .namespace("smithy_core.interfaces.retries", ".")
@@ -46,8 +40,22 @@ private AwsConfiguration() {}
                     .build())
             .documentation(
                     "The retry strategy or options for configuring retry behavior. Can be either a configured RetryStrategy or RetryStrategyOptions to create one.")
-            .nullable(false)
-            .useDescriptor(true)
+            .build();
+
+    /**
+     * AWS-specific metadata for descriptor-based config properties.
+     */
+    public static final AwsConfigPropertyMetadata REGION_METADATA = AwsConfigPropertyMetadata.builder()
+            .validator(Symbol.builder()
+                    .name("validate_region")
+                    .namespace("smithy_aws_core.config.validators", ".")
+                    .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                    .build())
+            .build();
+    /**
+     * AWS-specific metadata for descriptor-based config properties.
+     */
+    public static final AwsConfigPropertyMetadata RETRY_STRATEGY_METADATA = AwsConfigPropertyMetadata.builder()
             .validator(Symbol.builder()
                     .name("validate_retry_strategy")
                     .namespace("smithy_aws_core.config.validators", ".")

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsUserAgentIntegration.java

@@ -21,6 +21,14 @@
 @SmithyInternalApi
 public class AwsUserAgentIntegration implements PythonIntegration {
 
+    public static final AwsConfigPropertyMetadata SDK_UA_APP_ID_METADATA = AwsConfigPropertyMetadata.builder()
+            .validator(Symbol.builder()
+                    .name("validate_ua_string")
+                    .namespace("smithy_aws_core.config.validators", ".")
+                    .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
+                    .build())
+            .build();
+
     public static final String USER_AGENT_PLUGIN = """
             def aws_user_agent_plugin(config: $1T):
                 config.interceptors.append(
@@ -49,12 +57,6 @@ public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
                             "A unique and opaque application ID that is appended to the User-Agent header.")
                     .type(Symbol.builder().name("str").build())
                     .nullable(true)
-                    .useDescriptor(true)
-                    .validator(Symbol.builder()
-                            .name("validate_ua_string")
-                            .namespace("smithy_aws_core.config.validators", ".")
-                            .addDependency(AwsPythonDependency.SMITHY_AWS_CORE)
-                            .build())
                     .build();
 
             final String user_agent_plugin_file = "user_agent";