replace importing the wf with ast parsing #516
Conversation
Signed-off-by: Ayush Kamat <[email protected]>
|
|
||
|
|
||
| def is_task_decorator(decorator_name: str) -> bool: | ||
| return decorator_name in { |
There was a problem hiding this comment.
id just store the set as a global constant and do the in inline
|
|
||
| elif isinstance(decorator, ast.Call): | ||
| func = decorator.func | ||
| assert isinstance(func, ast.Name) |
There was a problem hiding this comment.
might want to handle ast.Attribute here if people import qualified e.g.
from latch import tasks
@tasks.small_task
def f():
....There was a problem hiding this comment.
ideally we would also try to resolve to a fully qualified name (i.e. only match latch.tasks.small_task vs any decorator that happens to share that name)
I'd leave a todo
| continue | ||
|
|
||
| if func.id == "workflow": | ||
| self.flyte_objects.append(FlyteObject("workflow", node.name)) |
There was a problem hiding this comment.
you probably want to store a qualified name? i.e. including the module names? flyte uses that iirc e.g. a workflow might be named src.main.wf because def wf() is in src/main.py
| continue | ||
|
|
||
| dockerfile: Optional[Path] = None | ||
| for kw in decorator.keywords: |
There was a problem hiding this comment.
double-check that dockerfile is a keyword-only argument?
id leave a note here
| continue | ||
|
|
||
| try: | ||
| dockerfile = Path(ast.literal_eval(kw.value)).resolve() |
There was a problem hiding this comment.
careful with the working directory here
if the path is relative you need to resolve it relative to project root probably? or the file that is being parsed? unclear rn
|
|
||
|
|
||
| def get_flyte_objects(file: Path) -> list[FlyteObject]: | ||
| res = [] |
| try: | ||
| parsed = ast.parse(file.read_text(), filename=file) | ||
| except SyntaxError as e: | ||
| click.secho(f"There is a syntax error in {file}: {e}", fg="red") |
There was a problem hiding this comment.
| click.secho(f"There is a syntax error in {file}: {e}", fg="red") | |
| traceback.print_exc() | |
| click.secho(f"\nRegistration failed due to a syntax error (see above)", fg="red") |
| self.file = file | ||
| self.flyte_objects: list[FlyteObject] = [] | ||
|
|
||
| def visit_FunctionDef(self, node: ast.FunctionDef): |
There was a problem hiding this comment.
you might want to skip functions that are not at the top level? what does flyte do with e.g. class methods?
| res = [] | ||
| if file.is_dir(): | ||
| for child in file.iterdir(): | ||
| res.extend(get_flyte_objects(child)) |
There was a problem hiding this comment.
might want to use a queue instead of recursing
the file tree might be very deep for some stupid reason
also might want to avoid recursing into directories that don't have a __init__.py since those cannot contain importable python files anyway
| return res | ||
|
|
||
| assert file.is_file() | ||
| if file.suffix != ".py": |
There was a problem hiding this comment.
not a great heuristic tbh
would be better to follow the import graph i think
leave a todo at least
_import_flyte_objectsbc that is used elsewhere and i dont have the time to fix those places right now