handle case when pandera is run with optimized python mode (#1749)

Signed-off-by: cosmicBboy <[email protected]>

Author: cosmicBboy

pandera/api/extensions.py

@@ -14,13 +14,6 @@
 from pandera.api.hypotheses import Hypothesis
 from pandera.strategies.base_strategies import STRATEGY_DISPATCHER
 
-try:
-    import pyspark.sql as ps
-
-    PYSPARK_INSTALLED = True
-except ImportError:  # pragma: no cover
-    PYSPARK_INSTALLED = False
-
 
 class BuiltinCheckRegistrationError(Exception):
     """
@@ -180,9 +173,17 @@ def register_check_method(  # pylint:disable=too-many-branches
     :return: register check function wrapper.
     """
 
-    # pylint: disable=import-outside-toplevel
+    # pylint: disable=import-outside-toplevel,too-many-statements
     from pandera.strategies.pandas_strategies import register_check_strategy
 
+    # NOTE: this needs to handle different dataframe types more elegantly
+    try:
+        import pyspark.sql as ps
+
+        PYSPARK_INSTALLED = True
+    except ImportError:  # pragma: no cover
+        PYSPARK_INSTALLED = False
+
     if statistics is None:
         statistics = []
 

pandera/engines/pyspark_engine.py

@@ -17,19 +17,15 @@
 
 import pyspark
 import pyspark.sql.types as pst
+from pyspark.sql import DataFrame
 from packaging import version
 
 from pandera import dtypes, errors
 from pandera.dtypes import immutable
 from pandera.engines import engine
-from pandera.engines.type_aliases import PysparkObject
 
-try:
-    import pyarrow  # pylint:disable=unused-import
 
-    PYARROW_INSTALLED = True
-except ImportError:  # pragma: no cover
-    PYARROW_INSTALLED = False
+PysparkObject = Union[DataFrame]
 
 
 DEFAULT_PYSPARK_PREC = pst.DecimalType().precision

pandera/engines/type_aliases.py

@@ -5,16 +5,7 @@
 import numpy as np
 import pandas as pd
 
-try:
-    from pyspark.sql import DataFrame
-
-    PYSPARK_INSTALLED = True
-except ImportError:  # pragma: no cover
-    PYSPARK_INSTALLED = False
 
 PandasObject = Union[pd.Series, pd.DataFrame]
 PandasExtensionType = pd.core.dtypes.base.ExtensionDtype
 PandasDataType = Union[pd.core.dtypes.base.ExtensionDtype, np.dtype, type]
-
-if PYSPARK_INSTALLED:
-    PysparkObject = Union[DataFrame]

pandera/utils.py

@@ -9,6 +9,11 @@ def docstring_substitution(*args: Any, **kwargs: Any) -> Callable[[F], F]:
     """Typed wrapper around pandas.util.Substitution."""
 
     def decorator(func: F) -> F:
+        # handle case when pandera is run in optimized mode:
+        # https://docs.python.org/3/using/cmdline.html#cmdoption-OO
+        if func.__doc__ is None:
+            return func
+
         if args:
             _doc = func.__doc__ % tuple(args)  # type: ignore[operator]
         elif kwargs: