[airflow] branch main updated: Add XComArg to lazy-imported list of Airflow module (#22862)

[as...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ash 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 90cc4b2b36 Add XComArg to lazy-imported list of Airflow module (#22862)
90cc4b2b36 is described below

commit 90cc4b2b361ad2b76c07be615fcb9434b932b02a
Author: Ash Berlin-Taylor <as...@firemirror.com>
AuthorDate: Fri Apr 8 17:12:49 2022 +0100

    Add XComArg to lazy-imported list of Airflow module (#22862)
    
    In writing the docs for Dynamic Task Mapping (AIP-42) I noticed that
    there are some cases where users need to use XComArg directly, and it
    didn't feel right to make the import things from `airflow.models`.
    
    And I've now refactored the lazy import to be "data-driven" as three
    blocks of almost identical code was my limit.
    
    Co-authored-by: Jed Cunningham <66...@users.noreply.github.com>
    Co-authored-by: D. Ferruzzi <fe...@amazon.com>
---
 airflow/__init__.py | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/airflow/__init__.py b/airflow/__init__.py
index 342122093f..93377b2df1 100644
--- a/airflow/__init__.py
+++ b/airflow/__init__.py
@@ -37,7 +37,7 @@ from airflow import version
 
 __version__ = version.version
 
-__all__ = ['__version__', 'login', 'DAG', 'PY36', 'PY37', 'PY38', 'PY39', 'PY310']
+__all__ = ['__version__', 'login', 'DAG', 'PY36', 'PY37', 'PY38', 'PY39', 'PY310', 'XComArg']
 
 # Make `airflow` an namespace package, supporting installing
 # airflow.providers.* in different locations (i.e. one in site, and one in user
@@ -54,18 +54,31 @@ PY38 = sys.version_info >= (3, 8)
 PY39 = sys.version_info >= (3, 9)
 PY310 = sys.version_info >= (3, 10)
 
+# Things to lazy import in form 'name': 'path.to.module'
+__lazy_imports = {
+    'DAG': 'airflow.models.dag',
+    'XComArg': 'airflow.models.xcom_arg',
+    'AirflowException': 'airflow.exceptions',
+}
+
 
 def __getattr__(name):
     # PEP-562: Lazy loaded attributes on python modules
-    if name == "DAG":
-        from airflow.models.dag import DAG
+    path = __lazy_imports.get(name)
+    if not path:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    import operator
 
-        return DAG
-    if name == "AirflowException":
-        from airflow.exceptions import AirflowException
+    # Strip off the "airflow." prefix because of how `__import__` works (it always returns the top level
+    # module)
+    without_prefix = path.split('.', 1)[-1]
 
-        return AirflowException
-    raise AttributeError(f"module {__name__} has no attribute {name}")
+    getter = operator.attrgetter(f'{without_prefix}.{name}')
+    val = getter(__import__(path))
+    # Store for next time
+    globals()[name] = val
+    return val
 
 
 if not settings.LAZY_LOAD_PLUGINS:
@@ -89,6 +102,7 @@ STATICA_HACK = True
 globals()['kcah_acitats'[::-1].upper()] = False
 if STATICA_HACK:  # pragma: no cover
     from airflow.models.dag import DAG
+    from airflow.models.xcom_arg import XComArg
     from airflow.exceptions import AirflowException