[GitHub] ashb closed pull request #3571: [AIRFLOW-2099] Handle getsource() calls gracefully

[GitBox <gi...@apache.org>]

ashb closed pull request #3571: [AIRFLOW-2099] Handle getsource() calls gracefully
URL: https://github.com/apache/incubator-airflow/pull/3571
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/airflow/www/utils.py b/airflow/www/utils.py
index 44fa5c4dcd..1bbc0936b3 100644
--- a/airflow/www/utils.py
+++ b/airflow/www/utils.py
@@ -17,6 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 #
+import inspect
 from future import standard_library
 standard_library.install_aliases()
 from builtins import str
@@ -375,6 +376,33 @@ def make_cache_key(*args, **kwargs):
     return (path + args).encode('ascii', 'ignore')
 
 
+def get_python_source(x):
+    """
+    Helper function to get Python source (or not), preventing exceptions
+    """
+    source_code = None
+
+    if isinstance(x, functools.partial):
+        source_code = inspect.getsource(x.func)
+
+    if source_code is None:
+        try:
+            source_code = inspect.getsource(x)
+        except TypeError:
+            pass
+
+    if source_code is None:
+        try:
+            source_code = inspect.getsource(x.__call__)
+        except (TypeError, AttributeError):
+            pass
+
+    if source_code is None:
+        source_code = 'No source code available for {}'.format(type(x))
+
+    return source_code
+
+
 class AceEditorWidget(wtforms.widgets.TextArea):
     """
     Renders an ACE code editor.
diff --git a/airflow/www/views.py b/airflow/www/views.py
index d37c0db45d..9a2fa5e8a9 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -249,7 +249,9 @@ def wrapped_markdown(s):
     'doc_yaml': lambda x: render(x, lexers.YamlLexer),
     'doc_md': wrapped_markdown,
     'python_callable': lambda x: render(
-        inspect.getsource(x), lexers.PythonLexer),
+        wwwutils.get_python_source(x),
+        lexers.PythonLexer,
+    ),
 }
 
 
diff --git a/tests/core.py b/tests/core.py
index fcbc0cfd3a..4e19687955 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -1842,6 +1842,20 @@ def test_fetch_task_instance(self):
         response = self.app.get(url)
         self.assertIn("print_the_context", response.data.decode('utf-8'))
 
+    def test_dag_view_task_with_python_operator_using_partial(self):
+        response = self.app.get(
+            '/admin/airflow/task?'
+            'task_id=test_dagrun_functool_partial&dag_id=test_task_view_type_check&'
+            'execution_date={}'.format(DEFAULT_DATE_DS))
+        self.assertIn("A function with two args", response.data.decode('utf-8'))
+
+    def test_dag_view_task_with_python_operator_using_instance(self):
+        response = self.app.get(
+            '/admin/airflow/task?'
+            'task_id=test_dagrun_instance&dag_id=test_task_view_type_check&'
+            'execution_date={}'.format(DEFAULT_DATE_DS))
+        self.assertIn("A __call__ method", response.data.decode('utf-8'))
+
     def tearDown(self):
         configuration.conf.set("webserver", "expose_config", "False")
         self.dag_bash.clear(start_date=DEFAULT_DATE, end_date=timezone.utcnow())
diff --git a/tests/dags/test_task_view_type_check.py b/tests/dags/test_task_view_type_check.py
new file mode 100644
index 0000000000..eee8058137
--- /dev/null
+++ b/tests/dags/test_task_view_type_check.py
@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""
+DAG designed to test a PythonOperator that calls a functool.partial
+"""
+import functools
+import logging
+
+from datetime import datetime
+
+from airflow.models import DAG
+from airflow.operators.python_operator import PythonOperator
+
+DEFAULT_DATE = datetime(2016, 1, 1)
+default_args = dict(
+    start_date=DEFAULT_DATE,
+    owner='airflow')
+
+
+class CallableClass(object):
+    def __call__(self):
+        """ A __call__ method """
+        pass
+
+
+def a_function(arg_x, arg_y):
+    """ A function with two args """
+    pass
+
+
+partial_function = functools.partial(a_function, arg_x=1)
+class_instance = CallableClass()
+
+logging.info('class_instance type: {}'.format(type(class_instance)))
+
+dag = DAG(dag_id='test_task_view_type_check', default_args=default_args)
+
+dag_task1 = PythonOperator(
+    task_id='test_dagrun_functool_partial',
+    dag=dag,
+    python_callable=partial_function,
+)
+
+dag_task2 = PythonOperator(
+    task_id='test_dagrun_instance',
+    dag=dag,
+    python_callable=class_instance,
+)
diff --git a/tests/www/test_utils.py b/tests/www/test_utils.py
index 9d788e88f1..891298c0a9 100644
--- a/tests/www/test_utils.py
+++ b/tests/www/test_utils.py
@@ -17,6 +17,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import functools
 import mock
 import unittest
 from xml.dom import minidom
@@ -190,6 +191,42 @@ def some_func():
                 self.assertEqual(anonymous_username, kwargs['owner'])
                 mocked_session_instance.add.assert_called_once()
 
+    def test_get_python_source_from_method(self):
+        class AMockClass(object):
+            def a_method(self):
+                """ A method """
+                pass
+
+        mocked_class = AMockClass()
+
+        result = utils.get_python_source(mocked_class.a_method)
+        self.assertIn('A method', result)
+
+    def test_get_python_source_from_class(self):
+        class AMockClass(object):
+            def __call__(self):
+                """ A __call__ method """
+                pass
+
+        mocked_class = AMockClass()
+
+        result = utils.get_python_source(mocked_class)
+        self.assertIn('A __call__ method', result)
+
+    def test_get_python_source_from_partial_func(self):
+        def a_function(arg_x, arg_y):
+            """ A function with two args """
+            pass
+
+        partial_function = functools.partial(a_function, arg_x=1)
+
+        result = utils.get_python_source(partial_function)
+        self.assertIn('A function with two args', result)
+
+    def test_get_python_source_from_none(self):
+        result = utils.get_python_source(None)
+        self.assertIn('No source code available', result)
+
 
 if __name__ == '__main__':
     unittest.main()


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services