Add support to specify language name in PapermillOperator

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 @@ def __init__(
         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 @@ def __init__(
         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 @@ def execute(self, context: 'Context'):
                 progress_bar=False,
                 report_mode=True,
                 kernel_name=self.kernel_name,
+                language_name=self.language_name,
             )

tests/providers/papermill/operators/test_papermill.py

@@ -31,6 +31,7 @@ def test_execute(self, mock_papermill):
         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 @@ def test_execute(self, mock_papermill):
             parameters=parameters,
             task_id="papermill_operator_test",
             kernel_name=kernel_name,
+            language_name=language_name,
             dag=None,
         )
 
@@ -50,6 +52,7 @@ def test_execute(self, mock_papermill):
             out_nb,
             parameters=parameters,
             kernel_name=kernel_name,
+            language_name=language_name,
             progress_bar=False,
             report_mode=True,
         )
@@ -64,14 +67,16 @@ def test_render_template(self):
             output_nb="/tmp/out-{{ dag.dag_id }}.ipynb",
             parameters={"msgs": "dag id is {{ dag.dag_id }}!"},
             kernel_name="python3",
+            language_name="python",
             dag=dag,
         )
 
         ti = TaskInstance(operator, run_id="papermill_test")
         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