Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -72,10 +72,23 @@ private void generateService(PythonWriter writer) {
}

writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
var asyncConfigSymbol = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model());
writer.write("""
def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None):
def __init__(
self,
config: $1T | $6T | None = None,
plugins: list[$2T] | None = None,
):
$3C
self._config = config or $1T()
if isinstance(config, $6T):
self._config: $1T = config # type: ignore[assignment]
elif isinstance(config, $1T) or config is None:
self._config = config or $1T()
else:
raise $7T(
f"config must be $6L or $1L, got {type(config).__name__}. "
f"Use 'await $6L.resolve()' instead."
)

client_plugins: list[$2T] = [
$4C
Expand All @@ -92,7 +105,9 @@ def __init__(self, config: $1T | None = None, plugins: list[$2T] | None = None):
pluginSymbol,
writer.consumer(w -> writeConstructorDocs(w, serviceSymbol.getName())),
writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)),
RuntimeTypes.RETRY_STRATEGY_RESOLVER);
RuntimeTypes.RETRY_STRATEGY_RESOLVER,
asyncConfigSymbol,
RuntimeTypes.EXPECTATION_NOT_MET_ERROR);

var topDownIndex = TopDownIndex.of(model);
var eventStreamIndex = EventStreamIndex.of(model);
Expand Down Expand Up @@ -249,7 +264,9 @@ private void writeSharedOperationInit(
raise $2T("protocol and transport MUST be set on the config to make calls.")

retry_strategy = await self._retry_strategy_resolver.resolve_retry_strategy(
retry_strategy=config.retry_strategy
retry_strategy=config.retry_strategy,
retry_mode=getattr(config, "retry_mode", None),
max_attempts=getattr(config, "max_attempts", None),
)

pipeline = $3T(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.logging.Logger;
import software.amazon.smithy.codegen.core.CodegenException;
import software.amazon.smithy.codegen.core.Symbol;
import software.amazon.smithy.aws.traits.ServiceTrait;
import software.amazon.smithy.model.Model;
import software.amazon.smithy.model.knowledge.NullableIndex;
import software.amazon.smithy.model.node.Node;
Expand Down Expand Up @@ -87,6 +88,49 @@ public static Symbol getPluginSymbol(PythonSettings settings) {
.build();
}

/**
* Gets the async configuration object symbol for the service.
*
* <p>This is the new async-resolved config class that inherits from AsyncAwsConfig.
* Derives the name from the SDK ID (e.g., "Bedrock Runtime" becomes
* "AsyncBedrockRuntimeConfig"). Falls back to "AsyncConfig" for non-AWS services.
*
* @param settings The client settings.
* @param model The model containing the service shape.
* @return Returns the async config symbol.
*/
public static Symbol getAsyncConfigSymbol(PythonSettings settings, Model model) {
var service = settings.service(model);
var name = service.getTrait(ServiceTrait.class)
.map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Config")
.orElse("AsyncConfig");
return Symbol.builder()
.name(name)
.namespace(String.format("%s.config", settings.moduleName()), ".")
.definitionFile(String.format("./src/%s/config.py", settings.moduleName()))
.build();
}

/**
* Gets the async plugin type hint symbol for the service.
*
* @param settings The client settings.
* @param model The model containing the service shape.
* @return Returns the async plugin type hint symbol.
*/
public static Symbol getAsyncPluginSymbol(PythonSettings settings, Model model) {
var service = settings.service(model);
var name = service.getTrait(ServiceTrait.class)
.map(trait -> "Async" + StringUtils.capitalize(trait.getSdkId()).replace(" ", "") + "Plugin")
.orElse("AsyncPlugin");
return Symbol.builder()
.name(name)
.namespace(String.format("%s.config", settings.moduleName()), ".")
.definitionFile(String.format("./src/%s/config.py", settings.moduleName()))
.build();
}


/**
* Gets the service error symbol.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,23 @@ public void run() {
writer.write("$L: TypeAlias = Callable[[$T], None]", plugin.getName(), config);
writer.writeDocs("A callable that allows customizing the config object on each request.", context);
});

// Generate the async config subclass and its plugin type
var model = context.model();
var asyncConfig = CodegenUtils.getAsyncConfigSymbol(context.settings(), model);
var asyncPlugin = CodegenUtils.getAsyncPluginSymbol(context.settings(), model);
context.writerDelegator().useFileWriter(asyncConfig.getDefinitionFile(), asyncConfig.getNamespace(), writer -> {
generateAsyncConfig(context, writer, asyncConfig);

// Generate the async plugin type alias
writer.addStdlibImport("typing", "Callable");
writer.addStdlibImport("typing", "TypeAlias");
writer.write("");
writer.write("");
writer.write("$L: TypeAlias = Callable[[$L], None]", asyncPlugin.getName(), asyncConfig.getName());
writer.writeDocs(
"A callable that allows customizing the async config object on each request.", context);
});
}

private void writeInterceptorsType(PythonWriter writer) {
Expand Down Expand Up @@ -340,10 +357,16 @@ private void generateConfig(GenerationContext context, PythonWriter writer) {
writer.pushState(new ConfigSection(finalProperties));
writer.addLocallyDefinedSymbol(configSymbol);
writer.addStdlibImport("dataclasses", "dataclass");
writer.addStdlibImport("warnings");
var asyncConfigName = CodegenUtils.getAsyncConfigSymbol(context.settings(), context.model()).getName();
writer.write("""
@dataclass(init=False)
class $L:
\"""Configuration for $L.\"""
\"""Configuration for $L.

.. deprecated::
Use :class:`$L` with ``await $L.resolve()`` instead.
\"""

${C|}

Expand All @@ -352,12 +375,22 @@ def __init__(
*,
${C|}
):
warnings.warn(
"$L is deprecated, use $L.resolve() instead. "
"This class will be removed in a future version.",
DeprecationWarning,
stacklevel=2,
)
${C|}
""",
configSymbol.getName(),
serviceId,
asyncConfigName,
asyncConfigName,
writer.consumer(w -> writePropertyDeclarations(w, finalProperties)),
writer.consumer(w -> writeInitParams(w, finalProperties)),
configSymbol.getName(),
asyncConfigName,
writer.consumer(w -> initializeProperties(w, finalProperties)));
writer.popState();
}
Expand Down Expand Up @@ -385,6 +418,160 @@ private void initializeProperties(PythonWriter writer, Collection<ConfigProperty
}
}

/**
* Generates the async config subclass that inherits from AsyncAwsConfig.
*
* <p>This class uses the FieldSpec-based resolution pipeline and adds
* service-specific fields (endpoint_resolver, protocol, auth_schemes,
* auth_scheme_resolver) with their defaults derived from the Smithy model.
*/
private void generateAsyncConfig(GenerationContext context, PythonWriter writer, Symbol asyncConfigSymbol) {
var model = context.model();
var service = context.settings().service(model);
final String serviceId = service.getTrait(ServiceTrait.class)
.map(ServiceTrait::getSdkId)
.orElse(context.settings().service().getName());

// Import AsyncAwsConfig base class
writer.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE);
var asyncAwsConfigSymbol = Symbol.builder()
.name("AsyncAwsConfig")
.namespace("smithy_aws_core.config.aws_config", ".")
.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE)
.build();

// Import FieldSpec and ClassVar
var fieldSpecSymbol = Symbol.builder()
.name("FieldSpec")
.namespace("smithy_aws_core.config.types", ".")
.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE)
.build();
writer.addStdlibImport("typing", "ClassVar");
writer.addStdlibImport("typing", "Any");
writer.addStdlibImport("dataclasses", "dataclass");

writer.write("");
writer.write("");
writer.write("@dataclass(kw_only=True)");
writer.openBlock("class $L($T):", asyncConfigSymbol.getName(), asyncAwsConfigSymbol);
writer.write("\"\"\"$L configuration (async-resolved).\"\"\"", serviceId);
writer.write("");

// Write service-specific field declarations
writer.write("endpoint_resolver: $T | None = None", RuntimeTypes.ENDPOINT_RESOLVER);
writer.write("protocol: $T | None = None", Symbol.builder()
.name("ClientProtocol[Any, Any]")
.addReference(Symbol.builder()
.name("ClientProtocol")
.namespace("smithy_core.aio.interfaces", ".")
.addDependency(SmithyPythonDependency.SMITHY_CORE)
.build())
.build());
writer.write("auth_schemes: dict[$T, $T] | None = None",
RuntimeTypes.SHAPE_ID,
Symbol.builder()
.name("AuthScheme[Any, Any, Any, Any]")
.addReference(Symbol.builder()
.name("AuthScheme")
.namespace("smithy_core.aio.interfaces.auth", ".")
.addDependency(SmithyPythonDependency.SMITHY_CORE)
.build())
.build());
writer.write("auth_scheme_resolver: $T | None = None",
CodegenUtils.getHttpAuthSchemeResolverSymbol(context.settings()));
writer.write("");

// Write _FIELDS class variable with service-specific defaults
writer.openBlock("_FIELDS: ClassVar[dict[str, $T]] = {", fieldSpecSymbol);
writer.write("**$T._FIELDS,", asyncAwsConfigSymbol);

// endpoint_uri FieldSpec — overrides base class with service-aware resolver
var makeEndpointResolverSymbol = Symbol.builder()
.name("make_endpoint_uri_resolver")
.namespace("smithy_aws_core.config.resolvers", ".")
.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE)
.build();
var snakeCaseServiceId = serviceId.replace(" ", "_").toLowerCase();
writer.write("\"endpoint_uri\": $T(", fieldSpecSymbol);
writer.indent();
writer.write("default=None,");
writer.write("resolver=$T($S),", makeEndpointResolverSymbol, snakeCaseServiceId);
writer.dedent();
writer.write("),");

// endpoint_resolver FieldSpec
var endpointPrefix = service.getTrait(ServiceTrait.class)
.map(ServiceTrait::getEndpointPrefix)
.orElse(context.settings().service().getName());
var standardRegionalResolverSymbol = Symbol.builder()
.name("StandardRegionalEndpointsResolver")
.namespace("smithy_aws_core.endpoints.standard_regional", ".")
.addDependency(SmithyPythonDependency.SMITHY_AWS_CORE)
.build();
writer.write("\"endpoint_resolver\": $T(", fieldSpecSymbol);
writer.indent();
writer.write("default_factory=lambda: $T(endpoint_prefix=$S),",
standardRegionalResolverSymbol, endpointPrefix);
writer.dedent();
writer.write("),");

// protocol FieldSpec
writer.write("\"protocol\": $T(", fieldSpecSymbol);
writer.indent();
writer.write("default_factory=lambda: ${C|},",
writer.consumer(w -> context.protocolGenerator().initializeProtocol(context, w)));
writer.dedent();
writer.write("),");

// auth_schemes FieldSpec
writer.write("\"auth_schemes\": $T(", fieldSpecSymbol);
writer.indent();
writer.write("default_factory=lambda: ${C|},",
writer.consumer(w -> writeAsyncDefaultAuthSchemes(context, w)));
writer.dedent();
writer.write("),");

// auth_scheme_resolver FieldSpec
writer.write("\"auth_scheme_resolver\": $T(", fieldSpecSymbol);
writer.indent();
writer.write("default_factory=HTTPAuthSchemeResolver,");
writer.dedent();
writer.write("),");

// transport FieldSpec
writer.write("\"transport\": $T(", fieldSpecSymbol);
writer.indent();
if (usesHttp2(context)) {
writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("awscrt"));
writer.write("default_factory=lambda: $T(),", RuntimeTypes.AWS_CRT_HTTP_CLIENT);
} else {
writer.addDependency(SmithyPythonDependency.SMITHY_HTTP.withOptionalDependencies("aiohttp"));
writer.write("default_factory=lambda: $T(),", RuntimeTypes.AIOHTTP_CLIENT);
}
writer.dedent();
writer.write("),");

writer.closeBlock("}");
writer.closeBlock("");
}

private static void writeAsyncDefaultAuthSchemes(GenerationContext context, PythonWriter writer) {
var service = context.settings().service(context.model());
writer.openBlock("{");
for (PythonIntegration integration : context.integrations()) {
for (RuntimeClientPlugin plugin : integration.getClientPlugins(context)) {
if (plugin.matchesService(context.model(), service) && plugin.getAuthScheme().isPresent()) {
var scheme = plugin.getAuthScheme().get();
writer.write("$T($S): ${C|},",
RuntimeTypes.SHAPE_ID,
scheme.getAuthTrait(),
writer.consumer(w -> scheme.initializeScheme(context, writer, service)));
}
}
}
writer.closeBlock("}");
}

private static final class AddAuthHelper implements CodeInterceptor<ConfigSection, PythonWriter> {
@Override
public Class<ConfigSection> sectionType() {
Expand Down
Loading
Loading