Skip to content

Commit

Permalink
Python: Type-stub generation for SSDKs (#2149)
Browse files Browse the repository at this point in the history
* Initial Python stub generation

* Handle default values correctly

* Only generate `__init__` for classes that have constructor signatures

* Preserve doc comments

* Make context class generic

* Put type hint into a string to fix runtime error

* Run `mypy` on CI

* Use `make` to build Python SSDKs while generating diffs

* Escape Python types in Rust comments

* Only mark class methods with

* Sort imports to minimize diffs

* Add type annotations for `PySocket`

* Dont extend classes from `object` as every class already implicitly extended from `object`

* Use `vars` instead of `inspect.getmembers` to skip inherited members of a class

* Fix linting issues

* Add some tests for stubgen and refactor it

* Add type annotations to `PyMiddlewareException`

* Fix tests on Python 3.7

Python 3.7 doesn't support reading signatures from `__text_signature__`
for non-builtin functions (i.e. C/Rust functions). For testing we're using
regular Python syntax for defining signature.

* Provide default values for `typing.Optional[T]` types in type-stubs

* Update `is_fn_like` to cover more cases

* Remove `tools/smithy-rs-tool-common/`

* Make `DECORATORS` an array instead of a list

* Ignore missing type stub errors for `aiohttp`
  • Loading branch information
unexge authored and LukeMathWalker committed Feb 13, 2023
1 parent 7787cb0 commit c7da1ac
Show file tree
Hide file tree
Showing 24 changed files with 1,412 additions and 69 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,174 @@
/*
* 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("]", "\\]")
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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))
}
Expand All @@ -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.
Expand All @@ -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(),
)
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -239,19 +244,33 @@ 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].
##[new]
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)?;
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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;
Expand Down

0 comments on commit c7da1ac

Please sign in to comment.