Skip to content

Commit f42882d

Browse files
ubaskotaubaskota
committed
Refactor and move config resolution codegen logic from generic layer to AWS layer
1 parent fd6d81a commit f42882d

7 files changed

Lines changed: 415 additions & 322 deletions

File tree

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package software.amazon.smithy.python.aws.codegen;
6+
7+
import java.util.Optional;
8+
import software.amazon.smithy.codegen.core.Symbol;
9+
import software.amazon.smithy.utils.SmithyInternalApi;
10+
11+
/**
12+
* AWS-specific config resolution metadata for a config property.
13+
* Holds validators, custom resolvers, and default values.
14+
*/
15+
@SmithyInternalApi
16+
public record AwsConfigPropertyMetadata(
17+
Symbol validator,
18+
Symbol customResolver,
19+
String defaultValue
20+
) {
21+
public Optional<Symbol> validatorOpt() {
22+
return Optional.ofNullable(validator);
23+
}
24+
25+
public Optional<Symbol> customResolverOpt() {
26+
return Optional.ofNullable(customResolver);
27+
}
28+
29+
public Optional<String> defaultValueOpt() {
30+
return Optional.ofNullable(defaultValue);
31+
}
32+
33+
public static Builder builder() {
34+
return new Builder();
35+
}
36+
37+
public static final class Builder {
38+
private Symbol validator;
39+
private Symbol customResolver;
40+
private String defaultValue;
41+
42+
public Builder validator(Symbol validator) {
43+
this.validator = validator;
44+
return this;
45+
}
46+
47+
public Builder customResolver(Symbol customResolver) {
48+
this.customResolver = customResolver;
49+
return this;
50+
}
51+
52+
public Builder defaultValue(String defaultValue) {
53+
this.defaultValue = defaultValue;
54+
return this;
55+
}
56+
57+
public AwsConfigPropertyMetadata build() {
58+
return new AwsConfigPropertyMetadata(validator, customResolver, defaultValue);
59+
}
60+
}
61+
}

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

Lines changed: 293 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,293 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
package software.amazon.smithy.python.aws.codegen;
6+
7+
import java.util.ArrayList;
8+
import java.util.List;
9+
import java.util.Map;
10+
import java.util.Set;
11+
import java.util.stream.Collectors;
12+
import software.amazon.smithy.python.codegen.ConfigProperty;
13+
import software.amazon.smithy.python.codegen.GenerationContext;
14+
import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
15+
import software.amazon.smithy.python.codegen.sections.ConfigSection;
16+
import software.amazon.smithy.python.codegen.writer.PythonWriter;
17+
import software.amazon.smithy.utils.CodeInterceptor;
18+
import software.amazon.smithy.utils.CodeSection;
19+
import software.amazon.smithy.utils.SmithyInternalApi;
20+
21+
/**
22+
* Intercepts the generated Config class to add AWS-specific descriptor-based
23+
* config resolution, keeping the generic ConfigGenerator unchanged.
24+
*/
25+
@SmithyInternalApi
26+
public class AwsConfigResolutionIntegration implements PythonIntegration {
27+
28+
// Metadata for properties that use descriptors, keyed by property name.
29+
private static final Map<String, AwsConfigPropertyMetadata> DESCRIPTOR_METADATA = Map.of(
30+
"region", AwsConfiguration.REGION_METADATA,
31+
"retry_strategy", AwsConfiguration.RETRY_STRATEGY_METADATA,
32+
"sdk_ua_app_id", AwsUserAgentIntegration.SDK_UA_APP_ID_METADATA
33+
);
34+
35+
@Override
36+
public List<? extends CodeInterceptor<? extends CodeSection, PythonWriter>> interceptors(
37+
GenerationContext context
38+
) {
39+
return List.of(new ConfigResolutionInterceptor());
40+
}
41+
42+
private static final class ConfigResolutionInterceptor
43+
implements CodeInterceptor<ConfigSection, PythonWriter> {
44+
45+
@Override
46+
public Class<ConfigSection> sectionType() {
47+
return ConfigSection.class;
48+
}
49+
50+
@Override
51+
public void write(PythonWriter writer, String previousText, ConfigSection section) {
52+
// Find properties that have descriptor metadata registered
53+
List<ConfigProperty> descriptorProps = section.properties().stream()
54+
.filter(p -> DESCRIPTOR_METADATA.containsKey(p.name()))
55+
.collect(Collectors.toList());
56+
57+
if (descriptorProps.isEmpty()) {
58+
writer.write(previousText);
59+
return;
60+
}
61+
62+
Set<String> descriptorNames = descriptorProps.stream()
63+
.map(ConfigProperty::name)
64+
.collect(Collectors.toSet());
65+
66+
addImports(writer, descriptorProps);
67+
68+
String transformed = transformConfigClass(previousText, descriptorProps, descriptorNames);
69+
writer.write(transformed);
70+
71+
appendGetSourceMethod(writer);
72+
}
73+
74+
private void addImports(PythonWriter writer, List<ConfigProperty> descriptorProps) {
75+
writer.addImport("smithy_aws_core.config.property", "ConfigProperty");
76+
writer.addImport("smithy_aws_core.config.resolver", "ConfigResolver");
77+
writer.addImport("smithy_aws_core.config.sources", "EnvironmentSource");
78+
writer.addImport("smithy_aws_core.config.source_info", "SourceInfo");
79+
80+
for (ConfigProperty prop : descriptorProps) {
81+
AwsConfigPropertyMetadata meta = DESCRIPTOR_METADATA.get(prop.name());
82+
if (meta == null) {
83+
continue;
84+
}
85+
meta.validatorOpt().ifPresent(sym ->
86+
writer.addImport(sym.getNamespace(), sym.getName()));
87+
meta.customResolverOpt().ifPresent(sym ->
88+
writer.addImport(sym.getNamespace(), sym.getName()));
89+
meta.defaultValueOpt().ifPresent(val -> {
90+
if (val.contains("RetryStrategyOptions")) {
91+
writer.addImport("smithy_core.retries", "RetryStrategyOptions");
92+
}
93+
});
94+
}
95+
}
96+
97+
private void appendGetSourceMethod(PythonWriter writer) {
98+
writer.write("""
99+
100+
def get_source(self, key: str) -> SourceInfo | None:
101+
\"""Get the source that provided a configuration value.
102+
103+
Args:
104+
key: The configuration key (e.g., 'region', 'retry_strategy')
105+
106+
Returns:
107+
The source info (SimpleSource or ComplexSource),
108+
or None if the key hasn't been resolved yet.
109+
\"""
110+
cached = self.__dict__.get(f'_cache_{key}')
111+
return cached[1] if cached else None
112+
""");
113+
}
114+
115+
private String transformConfigClass(
116+
String previousText,
117+
List<ConfigProperty> descriptorProps,
118+
Set<String> descriptorNames
119+
) {
120+
String[] lines = previousText.split("\n", -1);
121+
List<String> result = new ArrayList<>();
122+
123+
boolean inClassBody = false;
124+
boolean inInit = false;
125+
boolean initParamsStarted = false;
126+
boolean initBodyStarted = false;
127+
boolean descriptorsInserted = false;
128+
boolean resolverInitInserted = false;
129+
int i = 0;
130+
131+
while (i < lines.length) {
132+
String line = lines[i];
133+
String trimmed = line.trim();
134+
135+
if (trimmed.startsWith("class Config")) {
136+
inClassBody = true;
137+
result.add(line);
138+
i++;
139+
continue;
140+
}
141+
142+
// Skip field declarations for descriptor properties
143+
if (inClassBody && !inInit) {
144+
boolean isDescriptorField = false;
145+
for (String name : descriptorNames) {
146+
if (trimmed.startsWith(name + ":") || trimmed.startsWith(name + " :")) {
147+
isDescriptorField = true;
148+
break;
149+
}
150+
}
151+
if (isDescriptorField) {
152+
i++;
153+
// Skip following docstring if present
154+
while (i < lines.length) {
155+
String nextTrimmed = lines[i].trim();
156+
if (nextTrimmed.startsWith("\"\"\"")) {
157+
if (nextTrimmed.endsWith("\"\"\"") && nextTrimmed.length() > 3) {
158+
i++;
159+
} else {
160+
i++;
161+
while (i < lines.length && !lines[i].trim().endsWith("\"\"\"")) {
162+
i++;
163+
}
164+
i++;
165+
}
166+
break;
167+
} else if (nextTrimmed.isEmpty()) {
168+
i++;
169+
} else {
170+
break;
171+
}
172+
}
173+
continue;
174+
}
175+
}
176+
177+
// Detect __init__ definition start
178+
if (trimmed.startsWith("def __init__(")) {
179+
inInit = true;
180+
initParamsStarted = true;
181+
182+
if (!descriptorsInserted) {
183+
insertDescriptorDeclarations(result, descriptorProps);
184+
descriptorsInserted = true;
185+
}
186+
187+
result.add(line);
188+
i++;
189+
continue;
190+
}
191+
192+
// Detect end of __init__ params
193+
if (initParamsStarted && !initBodyStarted && trimmed.equals("):")) {
194+
result.add(line);
195+
initBodyStarted = true;
196+
i++;
197+
continue;
198+
}
199+
200+
// Insert resolver initialization at start of __init__ body
201+
if (initBodyStarted && !resolverInitInserted) {
202+
if (!trimmed.isEmpty()) {
203+
insertResolverInitialization(result, descriptorNames);
204+
resolverInitInserted = true;
205+
}
206+
}
207+
208+
// Skip self.X = X for descriptor properties in __init__ body
209+
if (initBodyStarted) {
210+
boolean isDescriptorInit = false;
211+
for (String name : descriptorNames) {
212+
if (trimmed.equals("self." + name + " = " + name)) {
213+
isDescriptorInit = true;
214+
break;
215+
}
216+
}
217+
if (isDescriptorInit) {
218+
i++;
219+
continue;
220+
}
221+
}
222+
223+
result.add(line);
224+
i++;
225+
}
226+
227+
return String.join("\n", result);
228+
}
229+
230+
private void insertDescriptorDeclarations(List<String> result, List<ConfigProperty> descriptorProps) {
231+
result.add(" # Config properties using descriptors");
232+
result.add(" _descriptors = {");
233+
for (ConfigProperty prop : descriptorProps) {
234+
AwsConfigPropertyMetadata meta = DESCRIPTOR_METADATA.get(prop.name());
235+
StringBuilder sb = new StringBuilder();
236+
sb.append(" '").append(prop.name()).append("': ConfigProperty('")
237+
.append(prop.name()).append("'");
238+
if (meta != null) {
239+
meta.validatorOpt().ifPresent(sym ->
240+
sb.append(", validator=").append(sym.getName()));
241+
meta.customResolverOpt().ifPresent(sym ->
242+
sb.append(", resolver_func=").append(sym.getName()));
243+
meta.defaultValueOpt().ifPresent(val ->
244+
sb.append(", default_value=").append(val));
245+
}
246+
sb.append("),");
247+
result.add(sb.toString());
248+
}
249+
result.add(" }");
250+
result.add("");
251+
252+
// Add class-level descriptor assignments with type hints
253+
for (ConfigProperty prop : descriptorProps) {
254+
String typeHint = prop.type().getName();
255+
if (prop.isNullable() && !typeHint.endsWith("| None")) {
256+
typeHint = typeHint + " | None";
257+
}
258+
result.add(" " + prop.name() + ": " + typeHint +
259+
" = _descriptors['" + prop.name() + "'] # type: ignore[assignment]");
260+
261+
if (!prop.documentation().isEmpty()) {
262+
String doc = prop.documentation();
263+
if (doc.contains("\n")) {
264+
result.add(" \"\"\"");
265+
for (String docLine : doc.split("\n")) {
266+
result.add(" " + docLine);
267+
}
268+
result.add(" \"\"\"");
269+
} else {
270+
result.add(" \"\"\"" + doc + "\"\"\"");
271+
}
272+
}
273+
result.add("");
274+
}
275+
276+
// Add _resolver field declaration
277+
result.add(" _resolver: ConfigResolver");
278+
result.add("");
279+
}
280+
281+
private void insertResolverInitialization(List<String> result, Set<String> descriptorNames) {
282+
result.add(" # Set instance values for descriptor properties");
283+
result.add(" self._resolver = ConfigResolver(sources=[EnvironmentSource()])");
284+
result.add("");
285+
result.add(" # Only set if provided (not None) to allow resolution from sources");
286+
result.add(" for key in self.__class__._descriptors.keys():");
287+
result.add(" value = locals().get(key)");
288+
result.add(" if value is not None:");
289+
result.add(" setattr(self, key, value)");
290+
result.add("");
291+
}
292+
}
293+
}

0 commit comments

Comments
 (0)