Skip to content

Instantly share code, notes, and snippets.

@mvandermeulen
Forked from zzstoatzz/prefect_and_generics.py
Created September 20, 2024 00:09
Show Gist options
  • Save mvandermeulen/7a72e29dcee5ffaf6f4d0ab9fbc9df2e to your computer and use it in GitHub Desktop.
Save mvandermeulen/7a72e29dcee5ffaf6f4d0ab9fbc9df2e to your computer and use it in GitHub Desktop.
generics in prefect
from typing import Sequence, TypeVar, Generic
from pydantic import BaseModel, TypeAdapter
from prefect import flow, task
ResultType = TypeVar("ResultType")
class Result(BaseModel, Generic[ResultType]):
data: Sequence[ResultType]
@task
def query() -> Result[str]:
'''provide `str` as the otherwise generic type parameter for `Result`'''
return Result[str](data='1 2 3'.split())
@task
def process_result(raw_result: Result[str]) -> Result[int]:
'''provide `int` as the otherwise generic type parameter for `Result`'''
return TypeAdapter(Result[int]).validate_python(raw_result.model_dump())
@flow(log_prints=True)
def work():
future = query.submit()
processed_result = process_result(future) # type: ignore
print(processed_result)
assert isinstance(processed_result, Result)
assert processed_result.data == [1, 2, 3]
if __name__ == "__main__":
work()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment