forked from smithy-lang/smithy-rs
-
Notifications
You must be signed in to change notification settings - Fork 0
/
PythonApplicationGenerator.kt
405 lines (391 loc) · 17.2 KB
/
PythonApplicationGenerator.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
/*
* 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.generators
import software.amazon.smithy.model.shapes.OperationShape
import software.amazon.smithy.model.traits.DocumentationTrait
import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
import software.amazon.smithy.rust.codegen.core.rustlang.rust
import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext
import software.amazon.smithy.rust.codegen.core.smithy.ErrorsModule
import software.amazon.smithy.rust.codegen.core.smithy.InputsModule
import software.amazon.smithy.rust.codegen.core.smithy.OutputsModule
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
import software.amazon.smithy.rust.codegen.core.util.getTrait
import software.amazon.smithy.rust.codegen.core.util.inputShape
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.smithy.ServerCargoDependency
import software.amazon.smithy.rust.codegen.server.smithy.generators.protocol.ServerProtocol
/**
* Generates a Python compatible application and server that can be configured from Python.
*
* Example:
* from pool import DatabasePool
* from my_library import App, OperationInput, OperationOutput
* @dataclass
* class Context:
* db = DatabasePool()
*
* app = App()
* app.context(Context())
*
* @app.operation
* def operation(input: OperationInput, ctx: State) -> OperationOutput:
* description = await ctx.db.get_description(input.name)
* return OperationOutput(description)
*
* app.run()
*
* The application holds a mapping between operation names (lowercase, snakecase),
* the context as defined in Python and some task local with the Python event loop
* for the current process.
*
* The application exposes several methods to Python:
* * `App()`: constructor to create an instance of `App`.
* * `run()`: run the application on a number of workers.
* * `context()`: register the context object that is passed to the Python handlers.
* * One register method per operation that can be used as decorator. For example if
* the model has one operation called `RegisterServer`, it will codegenerate a method
* of `App` called `register_service()` that can be used to decorate the Python implementation
* of this operation.
*
* This class also renders the implementation of the `aws_smithy_http_server_python::PyServer` trait,
* that abstracts the processes / event loops / workers lifecycles.
*/
class PythonApplicationGenerator(
codegenContext: CodegenContext,
private val protocol: ServerProtocol,
private val operations: List<OperationShape>,
) {
private val symbolProvider = codegenContext.symbolProvider
private val libName = codegenContext.settings.moduleName.toSnakeCase()
private val runtimeConfig = codegenContext.runtimeConfig
private val service = codegenContext.serviceShape
private val serviceName = service.id.name.toPascalCase()
private val model = codegenContext.model
private val codegenScope =
arrayOf(
"SmithyPython" to PythonServerCargoDependency.smithyHttpServerPython(runtimeConfig).toType(),
"SmithyServer" to ServerCargoDependency.smithyHttpServer(runtimeConfig).toType(),
"pyo3" to PythonServerCargoDependency.PyO3.toType(),
"pyo3_asyncio" to PythonServerCargoDependency.PyO3Asyncio.toType(),
"tokio" to PythonServerCargoDependency.Tokio.toType(),
"tracing" to PythonServerCargoDependency.Tracing.toType(),
"tower" to PythonServerCargoDependency.Tower.toType(),
"tower_http" to PythonServerCargoDependency.TowerHttp.toType(),
"num_cpus" to PythonServerCargoDependency.NumCpus.toType(),
"hyper" to PythonServerCargoDependency.Hyper.toType(),
"HashMap" to RuntimeType.HashMap,
"parking_lot" to PythonServerCargoDependency.ParkingLot.toType(),
"http" to RuntimeType.Http,
)
fun render(writer: RustWriter) {
renderPyApplicationRustDocs(writer)
renderAppStruct(writer)
renderAppDefault(writer)
renderAppClone(writer)
renderPyAppTrait(writer)
renderPyMethods(writer)
}
fun renderAppStruct(writer: RustWriter) {
writer.rustTemplate(
"""
##[#{pyo3}::pyclass]
##[derive(Debug)]
pub struct App {
handlers: #{HashMap}<String, #{SmithyPython}::PyHandler>,
middlewares: Vec<#{SmithyPython}::PyMiddlewareHandler>,
context: Option<#{pyo3}::PyObject>,
workers: #{parking_lot}::Mutex<Vec<#{pyo3}::PyObject>>,
}
""",
*codegenScope,
)
}
private fun renderAppClone(writer: RustWriter) {
writer.rustTemplate(
"""
impl Clone for App {
fn clone(&self) -> Self {
Self {
handlers: self.handlers.clone(),
middlewares: self.middlewares.clone(),
context: self.context.clone(),
workers: #{parking_lot}::Mutex::new(vec![]),
}
}
}
""",
*codegenScope,
)
}
private fun renderAppDefault(writer: RustWriter) {
writer.rustTemplate(
"""
impl Default for App {
fn default() -> Self {
Self {
handlers: Default::default(),
middlewares: vec![],
context: None,
workers: #{parking_lot}::Mutex::new(vec![]),
}
}
}
""",
"Protocol" to protocol.markerStruct(),
*codegenScope,
)
}
private fun renderPyAppTrait(writer: RustWriter) {
writer.rustBlockTemplate(
"""
impl #{SmithyPython}::PyApp for App
""",
*codegenScope,
) {
rustTemplate(
"""
fn workers(&self) -> &#{parking_lot}::Mutex<Vec<#{pyo3}::PyObject>> {
&self.workers
}
fn context(&self) -> &Option<#{pyo3}::PyObject> {
&self.context
}
fn handlers(&mut self) -> &mut #{HashMap}<String, #{SmithyPython}::PyHandler> {
&mut self.handlers
}
""",
*codegenScope,
)
rustBlockTemplate(
"""
fn build_service(&mut self, event_loop: &#{pyo3}::PyAny) -> #{pyo3}::PyResult<
#{tower}::util::BoxCloneService<
#{http}::Request<#{SmithyServer}::body::Body>,
#{http}::Response<#{SmithyServer}::body::BoxBody>,
std::convert::Infallible
>
>
""",
*codegenScope,
) {
rustTemplate(
"""
let builder = crate::service::$serviceName::builder_without_plugins();
""",
*codegenScope,
)
for (operation in operations) {
val operationName = symbolProvider.toSymbol(operation).name
val name = operationName.toSnakeCase()
rustTemplate(
"""
let ${name}_locals = #{pyo3_asyncio}::TaskLocals::new(event_loop);
let handler = self.handlers.get("$name").expect("Python handler for operation `$name` not found").clone();
let builder = builder.$name(move |input, state| {
#{pyo3_asyncio}::tokio::scope(${name}_locals.clone(), crate::operation_handler::$name(input, state, handler.clone()))
});
""",
*codegenScope,
)
}
rustTemplate(
"""
let mut service = #{tower}::util::BoxCloneService::new(builder.build().expect("one or more operations do not have a registered handler; this is a bug in the Python code generator, please file a bug report under https://github.com/awslabs/smithy-rs/issues"));
{
use #{tower}::Layer;
#{tracing}::trace!("adding middlewares to rust python router");
let mut middlewares = self.middlewares.clone();
// Reverse the middlewares, so they run with same order as they defined
middlewares.reverse();
for handler in middlewares {
#{tracing}::trace!(name = &handler.name, "adding python middleware");
let locals = #{pyo3_asyncio}::TaskLocals::new(event_loop);
let layer = #{SmithyPython}::PyMiddlewareLayer::<#{Protocol}>::new(handler, locals);
service = #{tower}::util::BoxCloneService::new(layer.layer(service));
}
}
Ok(service)
""",
"Protocol" to protocol.markerStruct(),
*codegenScope,
)
}
}
}
private fun renderPyMethods(writer: RustWriter) {
writer.rustBlockTemplate(
"""
##[#{pyo3}::pymethods]
impl App
""",
*codegenScope,
) {
rustTemplate(
"""
/// Create a new [App].
##[new]
pub fn new() -> Self {
Self::default()
}
/// Register a context object that will be shared between handlers.
##[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.
##[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)?;
#{tracing}::trace!(
name = &handler.name,
is_coroutine = handler.is_coroutine,
"registering middleware function",
);
self.middlewares.push(handler);
Ok(())
}
/// Main entrypoint: start the server on multiple workers.
##[pyo3(text_signature = "(${'$'}self, address, port, backlog, workers, tls)")]
pub fn run(
&mut self,
py: #{pyo3}::Python,
address: Option<String>,
port: Option<i32>,
backlog: Option<i32>,
workers: Option<usize>,
tls: Option<#{SmithyPython}::tls::PyTlsConfig>,
) -> #{pyo3}::PyResult<()> {
use #{SmithyPython}::PyApp;
self.run_server(py, address, port, backlog, workers, tls)
}
/// Lambda entrypoint: start the server on Lambda.
##[pyo3(text_signature = "(${'$'}self)")]
pub fn run_lambda(
&mut self,
py: #{pyo3}::Python,
) -> #{pyo3}::PyResult<()> {
use #{SmithyPython}::PyApp;
self.run_lambda_handler(py)
}
/// Build the service and start a single worker.
##[pyo3(text_signature = "(${'$'}self, socket, worker_number, tls)")]
pub fn start_worker(
&mut self,
py: pyo3::Python,
socket: &pyo3::PyCell<#{SmithyPython}::PySocket>,
worker_number: isize,
tls: Option<#{SmithyPython}::tls::PyTlsConfig>,
) -> pyo3::PyResult<()> {
use #{SmithyPython}::PyApp;
let event_loop = self.configure_python_event_loop(py)?;
let service = self.build_and_configure_service(py, event_loop)?;
self.start_hyper_worker(py, socket, event_loop, service, worker_number, tls)
}
""",
*codegenScope,
)
operations.map { operation ->
val operationName = symbolProvider.toSymbol(operation).name
val name = operationName.toSnakeCase()
rustTemplate(
"""
/// Method to register `$name` Python implementation inside the handlers map.
/// It can be used as a function decorator in Python.
##[pyo3(text_signature = "(${'$'}self, func)")]
pub fn $name(&mut self, py: #{pyo3}::Python, func: #{pyo3}::PyObject) -> #{pyo3}::PyResult<()> {
use #{SmithyPython}::PyApp;
self.register_operation(py, "$name", func)
}
""",
*codegenScope,
)
}
}
}
private fun renderPyApplicationRustDocs(writer: RustWriter) {
writer.rust(
"""
##[allow(clippy::tabs_in_doc_comments)]
/// Main Python application, used to register operations and context and start multiple
/// workers on the same shared socket.
///
/// Operations can be registered using the application object as a decorator (`@app.operation_name`).
///
/// Here's a full example to get you started:
///
/// ```python
""".trimIndent(),
)
writer.rust(
"""
/// from $libName import ${InputsModule.name}
/// from $libName import ${OutputsModule.name}
""".trimIndent(),
)
if (operations.any { it.errors.isNotEmpty() }) {
writer.rust("""/// from $libName import ${ErrorsModule.name}""".trimIndent())
}
writer.rust(
"""
/// from $libName import middleware
/// from $libName import App
///
/// @dataclass
/// class Context:
/// counter: int = 0
///
/// app = App()
/// app.context(Context())
///
/// @app.request_middleware
/// def request_middleware(request: middleware::Request):
/// if request.get_header("x-amzn-id") != "secret":
/// raise middleware.MiddlewareException("Unsupported `x-amz-id` header", 401)
///
""".trimIndent(),
)
writer.operationImplementationStubs(operations)
writer.rust(
"""
///
/// app.run()
/// ```
///
/// Any of operations above can be written as well prepending the `async` keyword and
/// the Python application will automatically handle it and schedule it on the event loop for you.
""".trimIndent(),
)
}
private fun RustWriter.operationImplementationStubs(operations: List<OperationShape>) = rust(
operations.joinToString("\n///\n") {
val operationDocumentation = it.getTrait<DocumentationTrait>()?.value
val ret = if (!operationDocumentation.isNullOrBlank()) {
operationDocumentation.replace("#", "##").prependIndent("/// ## ") + "\n"
} else ""
ret +
"""
/// ${it.signature()}:
/// raise NotImplementedError
""".trimIndent()
},
)
/**
* Returns the function signature for an operation handler implementation. Used in the documentation.
*/
private fun OperationShape.signature(): String {
val inputSymbol = symbolProvider.toSymbol(inputShape(model))
val outputSymbol = symbolProvider.toSymbol(outputShape(model))
val inputT = "${InputsModule.name}::${inputSymbol.name}"
val outputT = "${OutputsModule.name}::${outputSymbol.name}"
val operationName = symbolProvider.toSymbol(this).name.toSnakeCase()
return "@app.$operationName\n/// def $operationName(input: $inputT, ctx: Context) -> $outputT"
}
}