pydanticflow.py

from metaflow import FlowSpec, Parameter, step, JSONType, Task, current
from datetime import datetime
from functools import wraps
from pydantic import BaseModel

class Config(BaseModel):
    id: int
    name = 'John Doe'

def validate(f):
    @wraps(f)
    def wrapper(self):
        for param in self._graph_info['parameters']:
            config = getattr(self, param['name'])
            user = Config(**config)
        f(self)
    return wrapper

class PydanticFlow(FlowSpec):

    user = Parameter("user",  
                     help="Configuration",
                     type=JSONType)

    @validate
    @step
    def start(self):
        pass
        self.next(self.end)

    @step
    def end(self):
        print("done!")

if __name__ == '__main__':
    PydanticFlow()