From 8081eee8fa16877ef72c6875c1c06e24a1668122 Mon Sep 17 00:00:00 2001 From: marov Date: Sun, 17 Oct 2021 23:46:59 -0700 Subject: [PATCH 01/14] WIP airflow export --- .vscode/launch.json | 13 +++++++++++++ lineapy/cli/cli.py | 20 +++++++++++++++++++- lineapy/instrumentation/tracer.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index afef4f1da..7a1e1c1ff 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -28,6 +28,19 @@ "tests/housing.py", ], }, + { + "name": "lineapy --airflow", + "type": "python", + "request": "launch", + "module": "lineapy.cli.cli", + "args": [ + "--slice", + "p value", + "--airflow", + "sliced_housing", + "tests/housing.py", + ], + }, { "name": "Python: Current File", "type": "python", diff --git a/lineapy/cli/cli.py b/lineapy/cli/cli.py index 41d9531d3..760f309e7 100644 --- a/lineapy/cli/cli.py +++ b/lineapy/cli/cli.py @@ -37,6 +37,12 @@ default=None, help="Requires --slice. Export the sliced code that {slice} depends on to {export_slice}.py", ) +@click.option( + "--export-slice-to-airflow-dag", + "--airflow", + default=None, + help="Requires --slice. Export the sliced code that {slice} depends on to an Airflow DAG {export_slice}.py", +) @click.option( "--print-source", help="Whether to print the source code", is_flag=True ) @@ -59,6 +65,7 @@ def linea_cli( mode, slice, export_slice, + export_slice_to_airflow_dag, print_source, print_graph, verbose, @@ -79,7 +86,7 @@ def linea_cli( tracer = Tracer(db, SessionType.SCRIPT) transform(code, file_name, tracer) - if slice and not export_slice: + if slice and not export_slice and not export_slice_to_airflow_dag: tree.add( rich.console.Group( f"Slice of {repr(slice)}", @@ -94,6 +101,17 @@ def linea_cli( full_code = tracer.sliced_func(slice, export_slice) pathlib.Path(f"{export_slice}.py").write_text(full_code) + if export_slice_to_airflow_dag: + if not slice: + print( + "Please specify --slice. It is required for --export-slice-to-airflow-dag" + ) + exit(1) + full_code = tracer.sliced_aiflow_dag( + slice, export_slice_to_airflow_dag + ) + pathlib.Path(f"{export_slice}.py").write_text(full_code) + tracer.db.close() if print_graph: graph_code = tracer.graph.print( diff --git a/lineapy/instrumentation/tracer.py b/lineapy/instrumentation/tracer.py index 216cb61bd..be7ce2e74 100644 --- a/lineapy/instrumentation/tracer.py +++ b/lineapy/instrumentation/tracer.py @@ -1,11 +1,16 @@ import logging +import pathlib from collections import defaultdict from dataclasses import InitVar, dataclass, field from datetime import datetime from functools import cached_property +from importlib import import_module from os import getcwd from typing import Dict, Optional +from airflow import DAG +from airflow.operators.python_operator import PythonOperator +from airflow.utils.dates import days_ago from black import FileMode, format_str from lineapy.constants import GET_ITEM, GETATTR @@ -155,6 +160,30 @@ def sliced_func(self, slice_name: str, func_name: str) -> str: full_code = format_str(full_code, mode=black_mode) return full_code + def sliced_aiflow_dag(self, slice_name: str, func_name: str) -> str: + _sliced_func_code = self.sliced_func(slice_name, func_name) + pathlib.Path(f"{func_name}.py").write_text(_sliced_func_code) + _sliced_func = getattr(import_module(func_name), func_name) + DEFAULT_ARGS = { + "start_date": days_ago(1), + "owner": "airflow", + "retries": 2, + } + dagargs = { + "default_args": DEFAULT_ARGS, + "schedule_interval": "0 7 * * *", + "catchup": False, + "max_active_runs": 1, + "concurrency": 5, + } + dag = DAG(func_name, **dagargs) + with dag as dag: + task_id = func_name + task = PythonOperator( + task_id=task_id, python_callable=func_name, dag=dag + ) + return dag + def session_artifacts(self) -> list[ArtifactORM]: return self.db.get_artifacts_for_session(self.session_context.id) From 5623ced05c3a586a7a87ce453e956667746d9ec0 Mon Sep 17 00:00:00 2001 From: marov Date: Mon, 18 Oct 2021 00:07:48 -0700 Subject: [PATCH 02/14] typo --- lineapy/instrumentation/tracer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lineapy/instrumentation/tracer.py b/lineapy/instrumentation/tracer.py index be7ce2e74..f4ad95c0d 100644 --- a/lineapy/instrumentation/tracer.py +++ b/lineapy/instrumentation/tracer.py @@ -180,7 +180,7 @@ def sliced_aiflow_dag(self, slice_name: str, func_name: str) -> str: with dag as dag: task_id = func_name task = PythonOperator( - task_id=task_id, python_callable=func_name, dag=dag + task_id=task_id, python_callable=_sliced_func, dag=dag ) return dag From 76c036dc73bffb77a872e8d48d44408b967ee89d Mon Sep 17 00:00:00 2001 From: marov Date: Mon, 18 Oct 2021 22:05:03 -0700 Subject: [PATCH 03/14] Refactor --- .gitignore | 3 ++ .vscode/launch.json | 2 +- .vscode/settings.json | 4 +- lineapy/cli/cli.py | 7 ++-- lineapy/instrumentation/tracer.py | 62 +++++++++++-------------------- lineapy/plugins/airflow.py | 46 +++++++++++++++++++++++ 6 files changed, 77 insertions(+), 47 deletions(-) create mode 100644 lineapy/plugins/airflow.py diff --git a/.gitignore b/.gitignore index e82e91bfc..42064553e 100644 --- a/.gitignore +++ b/.gitignore @@ -156,6 +156,9 @@ housing.csv # Result of --export-slice sliced_housing.py +# Result of --export-slice-to-airflow-dag +sliced_housing_dag.py + tracer tracer.pdf devcontainer.json diff --git a/.vscode/launch.json b/.vscode/launch.json index 7a1e1c1ff..a777b023a 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -37,7 +37,7 @@ "--slice", "p value", "--airflow", - "sliced_housing", + "sliced_housing_dag", "tests/housing.py", ], }, diff --git a/.vscode/settings.json b/.vscode/settings.json index 1ace67d8e..fe61b0f47 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,7 +1,7 @@ { "python.formatting.provider": "black", "python.linting.mypyEnabled": true, - "python.linting.enabled": true, + "python.linting.enabled": false, "python.testing.pytestEnabled": true, "python.linting.flake8Enabled": true, "editor.formatOnSave": true, @@ -10,4 +10,4 @@ "source.organizeImports": true } } -} +} \ No newline at end of file diff --git a/lineapy/cli/cli.py b/lineapy/cli/cli.py index 0c341d5f7..95728fcab 100644 --- a/lineapy/cli/cli.py +++ b/lineapy/cli/cli.py @@ -11,6 +11,7 @@ from lineapy.db.relational.db import RelationalLineaDB from lineapy.instrumentation.tracer import Tracer from lineapy.logging import configure_logging +from lineapy.plugins.airflow import sliced_aiflow_dag from lineapy.transformer.node_transformer import transform from lineapy.utils import prettify @@ -117,10 +118,10 @@ def linea_cli( "Please specify --slice. It is required for --export-slice-to-airflow-dag" ) exit(1) - full_code = tracer.sliced_aiflow_dag( - slice, export_slice_to_airflow_dag + full_code = sliced_aiflow_dag( + tracer, slice, export_slice_to_airflow_dag ) - pathlib.Path(f"{export_slice}.py").write_text(full_code) + pathlib.Path(f"{export_slice_to_airflow_dag}.py").write_text(full_code) tracer.db.close() if print_graph: diff --git a/lineapy/instrumentation/tracer.py b/lineapy/instrumentation/tracer.py index aa11e00e7..0feea57a3 100644 --- a/lineapy/instrumentation/tracer.py +++ b/lineapy/instrumentation/tracer.py @@ -1,16 +1,11 @@ import logging -import pathlib from collections import defaultdict from dataclasses import InitVar, dataclass, field from datetime import datetime from functools import cached_property -from importlib import import_module from os import getcwd from typing import Dict, Optional -from airflow import DAG -from airflow.operators.python_operator import PythonOperator -from airflow.utils.dates import days_ago from black import FileMode, format_str from lineapy.constants import GET_ITEM, GETATTR @@ -130,18 +125,9 @@ def artifacts(self) -> dict[str, str]: def sliced_func(self, slice_name: str, func_name: str) -> str: artifact = self.db.get_artifact_by_name(slice_name) - if not artifact.node: + artifact_var = self.slice_var_name(artifact) + if not artifact_var: return "Unable to extract the slice" - _line_no = artifact.node.lineno if artifact.node.lineno else 0 - artifact_line = str(artifact.node.source_code.code).split("\n")[ - _line_no - 1 - ] - _col_offset = ( - artifact.node.col_offset if artifact.node.col_offset else 0 - ) - if _col_offset < 3: - return "Unable to extract the slice" - artifact_name = artifact_line[: _col_offset - 3] slice_code = get_program_slice(self.graph, [artifact.id]) # We split the code in import and code blocks and join them to full code test import_block, code_block, main_block = split_code_blocks( @@ -151,7 +137,7 @@ def sliced_func(self, slice_name: str, func_name: str) -> str: import_block + "\n\n" + code_block - + f"\n\treturn {artifact_name}" + + f"\n\treturn {artifact_var}" + "\n\n" + main_block ) @@ -161,30 +147,6 @@ def sliced_func(self, slice_name: str, func_name: str) -> str: full_code = format_str(full_code, mode=black_mode) return full_code - def sliced_aiflow_dag(self, slice_name: str, func_name: str) -> str: - _sliced_func_code = self.sliced_func(slice_name, func_name) - pathlib.Path(f"{func_name}.py").write_text(_sliced_func_code) - _sliced_func = getattr(import_module(func_name), func_name) - DEFAULT_ARGS = { - "start_date": days_ago(1), - "owner": "airflow", - "retries": 2, - } - dagargs = { - "default_args": DEFAULT_ARGS, - "schedule_interval": "0 7 * * *", - "catchup": False, - "max_active_runs": 1, - "concurrency": 5, - } - dag = DAG(func_name, **dagargs) - with dag as dag: - task_id = func_name - task = PythonOperator( - task_id=task_id, python_callable=_sliced_func, dag=dag - ) - return dag - def session_artifacts(self) -> list[ArtifactORM]: return self.db.get_artifacts_for_session(self.session_context.id) @@ -192,6 +154,24 @@ def slice(self, name: str) -> str: artifact = self.db.get_artifact_by_name(name) return get_program_slice(self.graph, [artifact.id]) + def slice_var_name(self, artifact: ArtifactORM) -> str: + """ + Returns the variable name for the given artifact. + i.e. in lineapy.linea_publish(p, "p value") "p" is returned + """ + if not artifact.node: + return None + _line_no = artifact.node.lineno if artifact.node.lineno else 0 + artifact_line = str(artifact.node.source_code.code).split("\n")[ + _line_no - 1 + ] + _col_offset = ( + artifact.node.col_offset if artifact.node.col_offset else 0 + ) + if _col_offset < 3: + return None + return artifact_line[: _col_offset - 3] + def visualize( self, filename="tracer", diff --git a/lineapy/plugins/airflow.py b/lineapy/plugins/airflow.py new file mode 100644 index 000000000..937d6f433 --- /dev/null +++ b/lineapy/plugins/airflow.py @@ -0,0 +1,46 @@ +import logging +import os +import pathlib +import tempfile +from importlib import import_module + +from black import FileMode, format_str + +from lineapy.graph_reader.program_slice import ( + get_program_slice, + split_code_blocks, +) +from lineapy.instrumentation.tracer import Tracer + + +def sliced_aiflow_dag(tracer: Tracer, slice_name: str, func_name: str) -> str: + """ + Returns a an Airflow DAG of the sliced code. + + :param tracer: the tracer object. + :param slice_name: name of the artifacts to get the code slice for. + :return: string containing the code of the Airflow DAG running this slice + """ + artifact = tracer.db.get_artifact_by_name(slice_name) + artifact_var = tracer.slice_var_name(artifact) + if not artifact_var: + return "Unable to extract the slice" + slice_code = get_program_slice(tracer.graph, [artifact.id]) + # We split the code in import and code blocks and join them to full code test + import_block, code_block, main_block = split_code_blocks( + slice_code, func_name + ) + # TODO Add DAG specific blocks + full_code = ( + import_block + + "\n\n" + + code_block + + f"\n\treturn {artifact_var}" + + "\n\n" + + main_block + ) + # Black lint + black_mode = FileMode() + black_mode.line_length = 79 + full_code = format_str(full_code, mode=black_mode) + return full_code From c4e119eccb377ce9db219b19afb15b60c6aa656f Mon Sep 17 00:00:00 2001 From: marov Date: Mon, 18 Oct 2021 22:22:04 -0700 Subject: [PATCH 04/14] Optput a working DAG --- lineapy/plugins/airflow.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/lineapy/plugins/airflow.py b/lineapy/plugins/airflow.py index 937d6f433..4cdd18eb3 100644 --- a/lineapy/plugins/airflow.py +++ b/lineapy/plugins/airflow.py @@ -12,6 +12,28 @@ ) from lineapy.instrumentation.tracer import Tracer +AIRFLOW_IMPORTS_TEMPLATE = """ +from airflow import DAG +from airflow.utils.dates import days_ago +from airflow.operators.python_operator import PythonOperator +""" + +AIRFLOW_MAIN_TEMPLATE = """ +default_dag_args = {"owner": "airflow", "retries": 2, "start_date": days_ago(1)} + +dag = DAG( + dag_id="DAG_NAME_dag", + schedule_interval="*/15 * * * *", # Every 15 minutes + max_active_runs=1, + catchup=False, + default_args=default_dag_args, +) + +DAG_NAME = PythonOperator( + dag=dag, task_id=f"DAG_NAME_task", python_callable=DAG_NAME, +) +""" + def sliced_aiflow_dag(tracer: Tracer, slice_name: str, func_name: str) -> str: """ @@ -30,14 +52,15 @@ def sliced_aiflow_dag(tracer: Tracer, slice_name: str, func_name: str) -> str: import_block, code_block, main_block = split_code_blocks( slice_code, func_name ) - # TODO Add DAG specific blocks full_code = ( import_block + + "\n" + + AIRFLOW_IMPORTS_TEMPLATE + "\n\n" + code_block - + f"\n\treturn {artifact_var}" + + f"\n\tprint({artifact_var})" # TODO What to do with artifact_var in a DAG? + "\n\n" - + main_block + + AIRFLOW_MAIN_TEMPLATE.replace("DAG_NAME", func_name) ) # Black lint black_mode = FileMode() From ff8fb3f4a186609e088289fd577a49adb8e19f60 Mon Sep 17 00:00:00 2001 From: marov Date: Mon, 18 Oct 2021 22:34:27 -0700 Subject: [PATCH 05/14] Unused imports --- lineapy/plugins/airflow.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/lineapy/plugins/airflow.py b/lineapy/plugins/airflow.py index 4cdd18eb3..1d72a5a20 100644 --- a/lineapy/plugins/airflow.py +++ b/lineapy/plugins/airflow.py @@ -1,9 +1,3 @@ -import logging -import os -import pathlib -import tempfile -from importlib import import_module - from black import FileMode, format_str from lineapy.graph_reader.program_slice import ( From de89e7a58d645f6b2c12e9f541989cf6ac95d7f8 Mon Sep 17 00:00:00 2001 From: marov Date: Mon, 18 Oct 2021 22:38:41 -0700 Subject: [PATCH 06/14] Return type --- lineapy/instrumentation/tracer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lineapy/instrumentation/tracer.py b/lineapy/instrumentation/tracer.py index 0feea57a3..2e117f362 100644 --- a/lineapy/instrumentation/tracer.py +++ b/lineapy/instrumentation/tracer.py @@ -160,7 +160,7 @@ def slice_var_name(self, artifact: ArtifactORM) -> str: i.e. in lineapy.linea_publish(p, "p value") "p" is returned """ if not artifact.node: - return None + return "" _line_no = artifact.node.lineno if artifact.node.lineno else 0 artifact_line = str(artifact.node.source_code.code).split("\n")[ _line_no - 1 @@ -169,7 +169,7 @@ def slice_var_name(self, artifact: ArtifactORM) -> str: artifact.node.col_offset if artifact.node.col_offset else 0 ) if _col_offset < 3: - return None + return "" return artifact_line[: _col_offset - 3] def visualize( From 06bf5821dbcdfb925ca78eab903474df88028b37 Mon Sep 17 00:00:00 2001 From: marov Date: Tue, 19 Oct 2021 23:45:32 -0700 Subject: [PATCH 07/14] Add test and docs --- CONTRIBUTING.md | 13 +++++++++++++ setup.py | 1 + tests/test_airflow.py | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+) create mode 100644 tests/test_airflow.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 70960c536..0621e9276 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -84,6 +84,19 @@ jupyter nbconvert --to notebook --execute tests/test_notebook.ipynb --inplace -- Or you can open it in a notebook UI (JupyterLab, JupyterNotebook, VS Code, etc.) and re-run it manually +### Airflow + +Sliced code can be exported to an Airflow DAG using the following command: + +``` +lineapy tests/housing.py --slice "p value" --airflow sliced_housing_dag +``` +This creates a `sliced_housing_dag.py` file in the current dir. It can be executed with: + +``` +airflow db init +airflow dags test sliced_housing_dag_dag $(date '+%Y-%m-%d') -S . +``` ## Visual Graphs Sometimes it's helpful to see a visual representation of the graph diff --git a/setup.py b/setup.py index 165f310b7..044448926 100644 --- a/setup.py +++ b/setup.py @@ -86,6 +86,7 @@ def version(path): "coveralls", "seaborn", "graphviz", + "apache-airflow==2.2.0", ] }, include_package_data=True, diff --git a/tests/test_airflow.py b/tests/test_airflow.py new file mode 100644 index 000000000..0cb844d81 --- /dev/null +++ b/tests/test_airflow.py @@ -0,0 +1,35 @@ +import subprocess + + +def test_export_slice_housing_dag(): + """ + Verifies that the "--export-slice" CLI command is aliased to the `lienapy` executable + """ + subprocess.check_call( + [ + "lineapy", + "tests/housing.py", + "--slice", + "p value", + "--airflow", + "sliced_housing_dag", + ] + ) + subprocess.check_call( + [ + "airflow", + "db", + "init", + ] + ) + subprocess.check_call( + [ + "airflow", + "dags", + "test", + "sliced_housing_dag_dag", + "2020-10-19", + "-S", + ".", + ] + ) From 6448004f48d19361fee47ca025a9f2e60b3e3fda Mon Sep 17 00:00:00 2001 From: marov Date: Tue, 19 Oct 2021 23:58:11 -0700 Subject: [PATCH 08/14] Update docstring --- tests/test_airflow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_airflow.py b/tests/test_airflow.py index 0cb844d81..f6c762fa9 100644 --- a/tests/test_airflow.py +++ b/tests/test_airflow.py @@ -3,7 +3,7 @@ def test_export_slice_housing_dag(): """ - Verifies that the "--export-slice" CLI command is aliased to the `lienapy` executable + Verifies that the "--airflow" CLI command produces a working Airflow DAG """ subprocess.check_call( [ From fdbfc0ff952b9b2b25133a0837695952fc61f110 Mon Sep 17 00:00:00 2001 From: marov Date: Tue, 19 Oct 2021 23:59:20 -0700 Subject: [PATCH 09/14] Airflow without version --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 044448926..83a33d853 100644 --- a/setup.py +++ b/setup.py @@ -86,7 +86,7 @@ def version(path): "coveralls", "seaborn", "graphviz", - "apache-airflow==2.2.0", + "apache-airflow", ] }, include_package_data=True, From 186af91243f3f9b8c4bb7283f44fcf02b3623661 Mon Sep 17 00:00:00 2001 From: marov Date: Wed, 20 Oct 2021 00:06:14 -0700 Subject: [PATCH 10/14] no airflow install --- setup.py | 1 - 1 file changed, 1 deletion(-) diff --git a/setup.py b/setup.py index 83a33d853..165f310b7 100644 --- a/setup.py +++ b/setup.py @@ -86,7 +86,6 @@ def version(path): "coveralls", "seaborn", "graphviz", - "apache-airflow", ] }, include_package_data=True, From acae9849630854b085927b2daeeaef9fe7d7f160 Mon Sep 17 00:00:00 2001 From: marov Date: Wed, 20 Oct 2021 00:11:10 -0700 Subject: [PATCH 11/14] apache-airflow==2.1.3 --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 165f310b7..e26e6ed04 100644 --- a/setup.py +++ b/setup.py @@ -86,6 +86,7 @@ def version(path): "coveralls", "seaborn", "graphviz", + "apache-airflow==2.1.3", ] }, include_package_data=True, From 1dc2cdd4ea925c457e59d452c8f5a86e82d27023 Mon Sep 17 00:00:00 2001 From: marov Date: Wed, 20 Oct 2021 00:18:15 -0700 Subject: [PATCH 12/14] 2.0.1 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index e26e6ed04..336a60d1a 100644 --- a/setup.py +++ b/setup.py @@ -86,7 +86,7 @@ def version(path): "coveralls", "seaborn", "graphviz", - "apache-airflow==2.1.3", + "apache-airflow==2.0.1", ] }, include_package_data=True, From 0caa77f26e85ddb1ce81605d8fe59566e2d79504 Mon Sep 17 00:00:00 2001 From: marov Date: Wed, 20 Oct 2021 09:56:32 -0700 Subject: [PATCH 13/14] mypy==0.770 --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 336a60d1a..a0f4c206c 100644 --- a/setup.py +++ b/setup.py @@ -71,7 +71,7 @@ def version(path): "sklearn", "flake8", "syrupy==1.4.5", - "mypy", + "mypy==0.770", "isort", "pytest", "matplotlib", @@ -86,7 +86,7 @@ def version(path): "coveralls", "seaborn", "graphviz", - "apache-airflow==2.0.1", + "apache-airflow==2.2.0", ] }, include_package_data=True, From c0b27458bef0a4550613270570d9cbd35b253031 Mon Sep 17 00:00:00 2001 From: Shardul Sardesai Date: Wed, 20 Oct 2021 11:42:30 -0700 Subject: [PATCH 14/14] skipping mypy plugins (sqlalchemy and pydantic) for now since they are blocking airflow work. will sync up with Saul to see if this is acceptable or if we have a better solution. --- lineapy/db/relational/schema/relational.py | 18 +++++++++--------- pyproject.toml | 2 +- setup.py | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/lineapy/db/relational/schema/relational.py b/lineapy/db/relational/schema/relational.py index cf6a9b0a4..9204eb5f4 100644 --- a/lineapy/db/relational/schema/relational.py +++ b/lineapy/db/relational/schema/relational.py @@ -96,7 +96,7 @@ def process_result_value(self, value, dialect): return value -class SessionContextORM(Base): +class SessionContextORM(Base): # type: ignore __tablename__ = "session_context" id = Column(String, primary_key=True) environment_type = Column(Enum(SessionType)) @@ -108,7 +108,7 @@ class SessionContextORM(Base): execution_id = Column(String, ForeignKey("execution.id")) -class LibraryORM(Base): +class LibraryORM(Base): # type: ignore __tablename__ = "library" __table_args__ = ( UniqueConstraint( @@ -125,7 +125,7 @@ class LibraryORM(Base): path = Column(String) -class ArtifactORM(Base): +class ArtifactORM(Base): # type: ignore """ An artifact is a named pointer to a node. """ @@ -140,7 +140,7 @@ class ArtifactORM(Base): ) -class ExecutionORM(Base): +class ExecutionORM(Base): # type: ignore """ An execution represents one Python interpreter invocation of some number of nodes """ @@ -150,7 +150,7 @@ class ExecutionORM(Base): timestamp = Column(DateTime, nullable=True, default=datetime.utcnow) -class NodeValueORM(Base): +class NodeValueORM(Base): # type: ignore """ A node value represents the value of a node during some execution. @@ -170,7 +170,7 @@ class NodeValueORM(Base): end_time = Column(DateTime, nullable=True) -class BaseNodeORM(Base): +class BaseNodeORM(Base): # type: ignore """ node.source_code has a path value if node.session.environment_type == "script" otherwise the environment type is "jupyter" and it has a jupyter execution @@ -214,7 +214,7 @@ class BaseNodeORM(Base): } -class SourceCodeORM(Base): +class SourceCodeORM(Base): # type: ignore __tablename__ = "source_code" id = Column(String, primary_key=True) @@ -261,7 +261,7 @@ class ImportNodeORM(BaseNodeORM): # https://docs.sqlalchemy.org/en/14/orm/basic_relationships.html#association-object -class PositionalArgORM(Base): +class PositionalArgORM(Base): # type: ignore __tablename__ = "positional_arg" call_node_id: str = Column( ForeignKey("call_node.id"), primary_key=True, nullable=False @@ -273,7 +273,7 @@ class PositionalArgORM(Base): argument = relationship(BaseNodeORM, uselist=False) -class KeywordArgORM(Base): +class KeywordArgORM(Base): # type: ignore __tablename__ = "keyword_arg" call_node_id: str = Column( ForeignKey("call_node.id"), primary_key=True, nullable=False diff --git a/pyproject.toml b/pyproject.toml index eaf6a7033..01c09f83a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ relative_files = true # https://docs.sqlalchemy.org/en/14/orm/extensions/mypy.html # https://pydantic-docs.helpmanual.io/mypy_plugin/#enabling-the-plugin -plugins = ["sqlalchemy.ext.mypy.plugin", "pydantic.mypy"] +# plugins = ["sqlalchemy.ext.mypy.plugin", "pydantic.mypy"] # Enable function body type checking, even if function types are not annotated diff --git a/setup.py b/setup.py index a0f4c206c..044448926 100644 --- a/setup.py +++ b/setup.py @@ -71,7 +71,7 @@ def version(path): "sklearn", "flake8", "syrupy==1.4.5", - "mypy==0.770", + "mypy", "isort", "pytest", "matplotlib",