This is an automated email from the ASF dual-hosted git repository.
joshfell pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6962d8a355 Add support to specify language name in PapermillOperator
(#23916)
6962d8a355 is described below
commit 6962d8a3556999af2eec459c944417ddd6d2cfb3
Author: John Green <[email protected]>
AuthorDate: Fri May 27 21:00:29 2022 +0100
Add support to specify language name in PapermillOperator (#23916)
* Add support to specify language name in PapermillOperator
* Replace getattr() with simple attribute access
---
airflow/providers/papermill/operators/papermill.py | 5 ++++-
tests/providers/papermill/operators/test_papermill.py | 13 +++++++++----
2 files changed, 13 insertions(+), 5 deletions(-)
diff --git a/airflow/providers/papermill/operators/papermill.py
b/airflow/providers/papermill/operators/papermill.py
index 36e8539b1e..b79ac226c1 100644
--- a/airflow/providers/papermill/operators/papermill.py
+++ b/airflow/providers/papermill/operators/papermill.py
@@ -50,7 +50,7 @@ class PapermillOperator(BaseOperator):
supports_lineage = True
- template_fields: Sequence[str] = ('input_nb', 'output_nb', 'parameters',
'kernel_name')
+ template_fields: Sequence[str] = ('input_nb', 'output_nb', 'parameters',
'kernel_name', 'language_name')
def __init__(
self,
@@ -59,6 +59,7 @@ class PapermillOperator(BaseOperator):
output_nb: Optional[str] = None,
parameters: Optional[Dict] = None,
kernel_name: Optional[str] = None,
+ language_name: Optional[str] = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -67,6 +68,7 @@ class PapermillOperator(BaseOperator):
self.output_nb = output_nb
self.parameters = parameters
self.kernel_name = kernel_name
+ self.language_name = language_name
if input_nb:
self.inlets.append(NoteBook(url=input_nb,
parameters=self.parameters))
if output_nb:
@@ -84,4 +86,5 @@ class PapermillOperator(BaseOperator):
progress_bar=False,
report_mode=True,
kernel_name=self.kernel_name,
+ language_name=self.language_name,
)
diff --git a/tests/providers/papermill/operators/test_papermill.py
b/tests/providers/papermill/operators/test_papermill.py
index 24d44e386a..e635524d01 100644
--- a/tests/providers/papermill/operators/test_papermill.py
+++ b/tests/providers/papermill/operators/test_papermill.py
@@ -31,6 +31,7 @@ class TestPapermillOperator(unittest.TestCase):
in_nb = "/tmp/does_not_exist"
out_nb = "/tmp/will_not_exist"
kernel_name = "python3"
+ language_name = "python"
parameters = {"msg": "hello_world", "train": 1}
op = PapermillOperator(
@@ -39,6 +40,7 @@ class TestPapermillOperator(unittest.TestCase):
parameters=parameters,
task_id="papermill_operator_test",
kernel_name=kernel_name,
+ language_name=language_name,
dag=None,
)
@@ -50,6 +52,7 @@ class TestPapermillOperator(unittest.TestCase):
out_nb,
parameters=parameters,
kernel_name=kernel_name,
+ language_name=language_name,
progress_bar=False,
report_mode=True,
)
@@ -64,6 +67,7 @@ class TestPapermillOperator(unittest.TestCase):
output_nb="/tmp/out-{{ dag.dag_id }}.ipynb",
parameters={"msgs": "dag id is {{ dag.dag_id }}!"},
kernel_name="python3",
+ language_name="python",
dag=dag,
)
@@ -71,7 +75,8 @@ class TestPapermillOperator(unittest.TestCase):
ti.dag_run = DagRun(execution_date=DEFAULT_DATE)
ti.render_templates()
- assert "/tmp/test_render_template.ipynb" == getattr(operator,
'input_nb')
- assert '/tmp/out-test_render_template.ipynb' == getattr(operator,
'output_nb')
- assert {"msgs": "dag id is test_render_template!"} ==
getattr(operator, 'parameters')
- assert "python3" == getattr(operator, 'kernel_name')
+ assert "/tmp/test_render_template.ipynb" == operator.input_nb
+ assert '/tmp/out-test_render_template.ipynb' == operator.output_nb
+ assert {"msgs": "dag id is test_render_template!"} ==
operator.parameters
+ assert "python3" == operator.kernel_name
+ assert "python" == operator.language_name
