Python: Type-stub generation for SSDKs (#2149) (e84ef6c7) · Commits · Public Repositories / Smithy Rust

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt

+1 −0

Original line number	Diff line number	Diff line
		@@ -413,6 +413,7 @@ class RustWriter private constructor(
		fun factory(debugMode: Boolean): Factory<RustWriter> = Factory { fileName: String, namespace: String ->
		when {
		fileName.endsWith(".toml") -> RustWriter(fileName, namespace, "#", debugMode = debugMode)
		fileName.endsWith(".py") -> RustWriter(fileName, namespace, "#", debugMode = debugMode)
		fileName.endsWith(".md") -> rawWriter(fileName, debugMode = debugMode)
		fileName == "LICENSE" -> rawWriter(fileName, debugMode = debugMode)
		fileName.startsWith("tests/") -> RustWriter(

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonType.kt

0 → 100644

+174 −0

Original line number	Diff line number	Diff line
		/*
		* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
		* SPDX-License-Identifier: Apache-2.0
		*/

		package software.amazon.smithy.rust.codegen.server.python.smithy

		import software.amazon.smithy.rust.codegen.core.rustlang.RustType

		/**
		* A hierarchy of Python types handled by Smithy codegen.
		*
		* Mostly copied from [RustType] and modified for Python accordingly.
		*/
		sealed class PythonType {
		/**
		* A Python type that contains [member], another [PythonType].
		* Used to generically operate over shapes that contain other shape.
		*/
		sealed interface Container {
		val member: PythonType
		val namespace: String?
		val name: String
		}

		/**
		* Name refers to the top-level type for import purposes.
		*/
		abstract val name: String

		open val namespace: String? = null

		object None : PythonType() {
		override val name: String = "None"
		}

		object Bool : PythonType() {
		override val name: String = "bool"
		}

		object Int : PythonType() {
		override val name: String = "int"
		}

		object Float : PythonType() {
		override val name: String = "float"
		}

		object Str : PythonType() {
		override val name: String = "str"
		}

		object Any : PythonType() {
		override val name: String = "Any"
		override val namespace: String = "typing"
		}

		data class List(override val member: PythonType) : PythonType(), Container {
		override val name: String = "List"
		override val namespace: String = "typing"
		}

		data class Dict(val key: PythonType, override val member: PythonType) : PythonType(), Container {
		override val name: String = "Dict"
		override val namespace: String = "typing"
		}

		data class Set(override val member: PythonType) : PythonType(), Container {
		override val name: String = "Set"
		override val namespace: String = "typing"
		}

		data class Optional(override val member: PythonType) : PythonType(), Container {
		override val name: String = "Optional"
		override val namespace: String = "typing"
		}

		data class Awaitable(override val member: PythonType) : PythonType(), Container {
		override val name: String = "Awaitable"
		override val namespace: String = "typing"
		}

		data class Callable(val args: kotlin.collections.List<PythonType>, val rtype: PythonType) : PythonType() {
		override val name: String = "Callable"
		override val namespace: String = "typing"
		}

		data class Union(val args: kotlin.collections.List<PythonType>) : PythonType() {
		override val name: String = "Union"
		override val namespace: String = "typing"
		}

		data class Opaque(override val name: String, val rustNamespace: String? = null) : PythonType() {
		// Since Python doesn't have a something like Rust's `crate::` we are using a custom placeholder here
		// and in our stub generation script we will replace placeholder with the real root module name.
		private val pythonRootModulePlaceholder = "__root_module_name__"

		override val namespace: String? = rustNamespace?.split("::")?.joinToString(".") {
		when (it) {
		"crate" -> pythonRootModulePlaceholder
		// In Python, we expose submodules from `aws_smithy_http_server_python`
		// like `types`, `middleware`, `tls` etc. from `__root_module__name`
		"aws_smithy_http_server_python" -> pythonRootModulePlaceholder
		else -> it
		}
		}
		}
		}

		/**
		* Return corresponding [PythonType] for a [RustType].
		*/
		fun RustType.pythonType(): PythonType =
		when (this) {
		is RustType.Unit -> PythonType.None
		is RustType.Bool -> PythonType.Bool
		is RustType.Float -> PythonType.Float
		is RustType.Integer -> PythonType.Int
		is RustType.String -> PythonType.Str
		is RustType.Vec -> PythonType.List(this.member.pythonType())
		is RustType.Slice -> PythonType.List(this.member.pythonType())
		is RustType.HashMap -> PythonType.Dict(this.key.pythonType(), this.member.pythonType())
		is RustType.HashSet -> PythonType.Set(this.member.pythonType())
		is RustType.Reference -> this.member.pythonType()
		is RustType.Option -> PythonType.Optional(this.member.pythonType())
		is RustType.Box -> this.member.pythonType()
		is RustType.Dyn -> this.member.pythonType()
		is RustType.Opaque -> PythonType.Opaque(this.name, this.namespace)
		// TODO(Constraints): How to handle this?
		// Revisit as part of https://github.com/awslabs/smithy-rs/issues/2114
		is RustType.MaybeConstrained -> this.member.pythonType()
		}

		/**
		* Render this type, including references and generic parameters.
		* It generates something like `typing.Dict[String, String]`.
		*/
		fun PythonType.render(fullyQualified: Boolean = true): String {
		val namespace = if (fullyQualified) {
		this.namespace?.let { "$it." } ?: ""
		} else ""
		val base = when (this) {
		is PythonType.None -> this.name
		is PythonType.Bool -> this.name
		is PythonType.Float -> this.name
		is PythonType.Int -> this.name
		is PythonType.Str -> this.name
		is PythonType.Any -> this.name
		is PythonType.Opaque -> this.name
		is PythonType.List -> "${this.name}[${this.member.render(fullyQualified)}]"
		is PythonType.Dict -> "${this.name}[${this.key.render(fullyQualified)}, ${this.member.render(fullyQualified)}]"
		is PythonType.Set -> "${this.name}[${this.member.render(fullyQualified)}]"
		is PythonType.Awaitable -> "${this.name}[${this.member.render(fullyQualified)}]"
		is PythonType.Optional -> "${this.name}[${this.member.render(fullyQualified)}]"
		is PythonType.Callable -> {
		val args = this.args.joinToString(", ") { it.render(fullyQualified) }
		val rtype = this.rtype.render(fullyQualified)
		"${this.name}[[$args], $rtype]"
		}
		is PythonType.Union -> {
		val args = this.args.joinToString(", ") { it.render(fullyQualified) }
		"${this.name}[$args]"
		}
		}
		return "$namespace$base"
		}

		/**
		* Renders [PythonType] with proper escaping for Docstrings.
		*/
		fun PythonType.renderAsDocstring(): String =
		this.render()
		.replace("[", "\\[")
		.replace("]", "\\]")

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/customizations/PythonServerCodegenDecorator.kt

+64 −0

Original line number	Diff line number	Diff line
		@@ -98,6 +98,7 @@ class PubUsePythonTypesDecorator : ServerCodegenDecorator {
		/**
		* Generates `pyproject.toml` for the crate.
		* - Configures Maturin as the build system
		* - Configures Python source directory
		*/
		class PyProjectTomlDecorator : ServerCodegenDecorator {
		override val name: String = "PyProjectTomlDecorator"
		@@ -110,6 +111,11 @@ class PyProjectTomlDecorator : ServerCodegenDecorator {
		"requires" to listOfNotNull("maturin>=0.14,<0.15"),
		"build-backend" to "maturin",
		).toMap(),
		"tool" to listOfNotNull(
		"maturin" to listOfNotNull(
		"python-source" to "python",
		).toMap(),
		).toMap(),
		)
		writeWithNoFormatting(TomlWriter().write(config))
		}
		@@ -134,6 +140,60 @@ class PyO3ExtensionModuleDecorator : ServerCodegenDecorator {
		}
		}

		/**
		* Generates `__init__.py` for the Python source.
		*
		* This file allows Python module to be imported like:
		* ```
		* import pokemon_service_server_sdk
		* pokemon_service_server_sdk.App()
		* ```
		* instead of:
		* ```
		* from pokemon_service_server_sdk import pokemon_service_server_sdk
		* ```
		*/
		class InitPyDecorator : ServerCodegenDecorator {
		override val name: String = "InitPyDecorator"
		override val order: Byte = 0

		override fun extras(codegenContext: ServerCodegenContext, rustCrate: RustCrate) {
		val libName = codegenContext.settings.moduleName.toSnakeCase()

		rustCrate.withFile("python/$libName/__init__.py") {
		writeWithNoFormatting(
		"""
		from .$libName import *

		__doc__ = $libName.__doc__
		if hasattr($libName, "__all__"):
		__all__ = $libName.__all__
		""".trimIndent(),
		)
		}
		}
		}

		/**
		* Generates `py.typed` for the Python source.
		*
		* This marker file is required to be PEP 561 compliant stub package.
		* Type definitions will be ignored by `mypy` if the package is not PEP 561 compliant:
		* https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-library-stubs-or-py-typed-marker
		*/
		class PyTypedMarkerDecorator : ServerCodegenDecorator {
		override val name: String = "PyTypedMarkerDecorator"
		override val order: Byte = 0

		override fun extras(codegenContext: ServerCodegenContext, rustCrate: RustCrate) {
		val libName = codegenContext.settings.moduleName.toSnakeCase()

		rustCrate.withFile("python/$libName/py.typed") {
		writeWithNoFormatting("")
		}
		}
		}

		val DECORATORS = arrayOf(
		/**
		* Add the [InternalServerError] error to all operations.
		@@ -150,4 +210,8 @@ val DECORATORS = arrayOf(
		PyProjectTomlDecorator(),
		// Add PyO3 extension module feature.
		PyO3ExtensionModuleDecorator(),
		// Generate `__init__.py` for the Python source.
		InitPyDecorator(),
		// Generate `py.typed` for the Python source.
		PyTypedMarkerDecorator(),
		)

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonApplicationGenerator.kt

+54 −2

Original line number	Diff line number	Diff line
		@@ -22,6 +22,8 @@ import software.amazon.smithy.rust.codegen.core.util.outputShape
		import software.amazon.smithy.rust.codegen.core.util.toPascalCase
		import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
		import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerCargoDependency
		import software.amazon.smithy.rust.codegen.server.python.smithy.PythonType
		import software.amazon.smithy.rust.codegen.server.python.smithy.renderAsDocstring
		import software.amazon.smithy.rust.codegen.server.smithy.ServerCargoDependency
		import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.ServerProtocol

		@@ -103,6 +105,9 @@ class PythonApplicationGenerator(
		"""
		##[#{pyo3}::pyclass]
		##[derive(Debug)]
		/// :generic Ctx:
		/// :extends typing.Generic\[Ctx\]:
		/// :rtype None:
		pub struct App {
		handlers: #{HashMap}<String, #{SmithyPython}::PyHandler>,
		middlewares: Vec<#{SmithyPython}::PyMiddlewareHandler>,
		@@ -239,6 +244,12 @@ class PythonApplicationGenerator(
		""",
		*codegenScope,
		) {
		val middlewareRequest = PythonType.Opaque("Request", "crate::middleware")
		val middlewareResponse = PythonType.Opaque("Response", "crate::middleware")
		val middlewareNext = PythonType.Callable(listOf(middlewareRequest), PythonType.Awaitable(middlewareResponse))
		val middlewareFunc = PythonType.Callable(listOf(middlewareRequest, middlewareNext), PythonType.Awaitable(middlewareResponse))
		val tlsConfig = PythonType.Opaque("TlsConfig", "crate::tls")

		rustTemplate(
		"""
		/// Create a new [App].
		@@ -246,12 +257,20 @@ class PythonApplicationGenerator(
		pub fn new() -> Self {
		Self::default()
		}

		/// Register a context object that will be shared between handlers.
		///
		/// :param context Ctx:
		/// :rtype ${PythonType.None.renderAsDocstring()}:
		##[pyo3(text_signature = "(${'$'}self, context)")]
		pub fn context(&mut self, context: #{pyo3}::PyObject) {
		self.context = Some(context);
		}

		/// Register a Python function to be executed inside a Tower middleware layer.
		///
		/// :param func ${middlewareFunc.renderAsDocstring()}:
		/// :rtype ${PythonType.None.renderAsDocstring()}:
		##[pyo3(text_signature = "(${'$'}self, func)")]
		pub fn middleware(&mut self, py: #{pyo3}::Python, func: #{pyo3}::PyObject) -> #{pyo3}::PyResult<()> {
		let handler = #{SmithyPython}::PyMiddlewareHandler::new(py, func)?;
		@@ -263,8 +282,16 @@ class PythonApplicationGenerator(
		self.middlewares.push(handler);
		Ok(())
		}

		/// Main entrypoint: start the server on multiple workers.
		##[pyo3(text_signature = "(${'$'}self, address, port, backlog, workers, tls)")]
		///
		/// :param address ${PythonType.Optional(PythonType.Str).renderAsDocstring()}:
		/// :param port ${PythonType.Optional(PythonType.Int).renderAsDocstring()}:
		/// :param backlog ${PythonType.Optional(PythonType.Int).renderAsDocstring()}:
		/// :param workers ${PythonType.Optional(PythonType.Int).renderAsDocstring()}:
		/// :param tls ${PythonType.Optional(tlsConfig).renderAsDocstring()}:
		/// :rtype ${PythonType.None.renderAsDocstring()}:
		##[pyo3(text_signature = "(${'$'}self, address=None, port=None, backlog=None, workers=None, tls=None)")]
		pub fn run(
		&mut self,
		py: #{pyo3}::Python,
		@@ -277,7 +304,10 @@ class PythonApplicationGenerator(
		use #{SmithyPython}::PyApp;
		self.run_server(py, address, port, backlog, workers, tls)
		}

		/// Lambda entrypoint: start the server on Lambda.
		///
		/// :rtype ${PythonType.None.renderAsDocstring()}:
		##[pyo3(text_signature = "(${'$'}self)")]
		pub fn run_lambda(
		&mut self,
		@@ -286,8 +316,9 @@ class PythonApplicationGenerator(
		use #{SmithyPython}::PyApp;
		self.run_lambda_handler(py)
		}

		/// Build the service and start a single worker.
		##[pyo3(text_signature = "(${'$'}self, socket, worker_number, tls)")]
		##[pyo3(text_signature = "(${'$'}self, socket, worker_number, tls=None)")]
		pub fn start_worker(
		&mut self,
		py: pyo3::Python,
		@@ -306,10 +337,31 @@ class PythonApplicationGenerator(
		operations.map { operation ->
		val operationName = symbolProvider.toSymbol(operation).name
		val name = operationName.toSnakeCase()

		val input = PythonType.Opaque("${operationName}Input", "crate::input")
		val output = PythonType.Opaque("${operationName}Output", "crate::output")
		val context = PythonType.Opaque("Ctx")
		val returnType = PythonType.Union(listOf(output, PythonType.Awaitable(output)))
		val handler = PythonType.Union(
		listOf(
		PythonType.Callable(
		listOf(input, context),
		returnType,
		),
		PythonType.Callable(
		listOf(input),
		returnType,
		),
		),
		)

		rustTemplate(
		"""
		/// Method to register `$name` Python implementation inside the handlers map.
		/// It can be used as a function decorator in Python.
		///
		/// :param func ${handler.renderAsDocstring()}:
		/// :rtype ${PythonType.None.renderAsDocstring()}:
		##[pyo3(text_signature = "(${'$'}self, func)")]
		pub fn $name(&mut self, py: #{pyo3}::Python, func: #{pyo3}::PyObject) -> #{pyo3}::PyResult<()> {
		use #{SmithyPython}::PyApp;

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerModuleGenerator.kt

+5 −0

Original line number	Diff line number	Diff line
		@@ -99,6 +99,7 @@ class PythonServerModuleGenerator(
		let types = #{pyo3}::types::PyModule::new(py, "types")?;
		types.add_class::<#{SmithyPython}::types::Blob>()?;
		types.add_class::<#{SmithyPython}::types::DateTime>()?;
		types.add_class::<#{SmithyPython}::types::Format>()?;
		types.add_class::<#{SmithyPython}::types::ByteStream>()?;
		#{pyo3}::py_run!(
		py,
		@@ -185,6 +186,10 @@ class PythonServerModuleGenerator(
		"""
		let aws_lambda = #{pyo3}::types::PyModule::new(py, "aws_lambda")?;
		aws_lambda.add_class::<#{SmithyPython}::lambda::PyLambdaContext>()?;
		aws_lambda.add_class::<#{SmithyPython}::lambda::PyClientApplication>()?;
		aws_lambda.add_class::<#{SmithyPython}::lambda::PyClientContext>()?;
		aws_lambda.add_class::<#{SmithyPython}::lambda::PyCognitoIdentity>()?;
		aws_lambda.add_class::<#{SmithyPython}::lambda::PyConfig>()?;
		pyo3::py_run!(
		py,
		aws_lambda,