fix the typing of conditionally importing xgboost and pai4sk

shinnar · Dec 3, 2020 · bd37c07 · bd37c07
1 parent ff1a1fb
commit bd37c07
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 5 deletions.
diff --git a/lale/lib/pai4sk/decision_tree_classifier.py b/lale/lib/pai4sk/decision_tree_classifier.py
@@ -11,13 +11,17 @@
 # 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.
+from typing import TYPE_CHECKING
 
 try:
-    import pai4sk
+    import pai4sk  # type: ignore
 
     pai4sk_installed = True
 except ImportError:
     pai4sk_installed = False
+    if TYPE_CHECKING:
+        import pai4sk  # type: ignore
+
 import lale.datasets.data_schemas
 import lale.docstrings
 import lale.operators

diff --git a/lale/lib/pai4sk/random_forest_classifier.py b/lale/lib/pai4sk/random_forest_classifier.py
@@ -12,12 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+
 try:
-    import pai4sk
+    import pai4sk  # type: ignore
 
     pai4sk_installed = True
 except ImportError:
     pai4sk_installed = False
+    if TYPE_CHECKING:
+        import pai4sk  # type: ignore
+
 import lale.datasets.data_schemas
 import lale.docstrings
 import lale.operators

diff --git a/lale/lib/pai4sk/random_forest_regressor.py b/lale/lib/pai4sk/random_forest_regressor.py
@@ -12,12 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+
 try:
-    import pai4sk
+    import pai4sk  # type: ignore
 
     pai4sk_installed = True
 except ImportError:
     pai4sk_installed = False
+    if TYPE_CHECKING:
+        import pai4sk  # type: ignore
+
+
 import lale.datasets.data_schemas
 import lale.docstrings
 import lale.operators

diff --git a/lale/lib/xgboost/xgb_classifier.py b/lale/lib/xgboost/xgb_classifier.py
@@ -12,18 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+
 import pandas as pd
 from sklearn.base import BaseEstimator
 
 import lale.docstrings
 import lale.operators
 
 try:
-    import xgboost
+    import xgboost  # type: ignore
 
     xgboost_installed = True
 except ImportError:
     xgboost_installed = False
+    if TYPE_CHECKING:
+        import xgboost  # type: ignore
 
 
 def _rename_one_feature(name):

diff --git a/lale/lib/xgboost/xgb_regressor.py b/lale/lib/xgboost/xgb_regressor.py
@@ -12,18 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import TYPE_CHECKING
+
 import pandas as pd
 from sklearn.base import BaseEstimator
 
 import lale.docstrings
 import lale.operators
 
 try:
-    import xgboost
+    import xgboost  # type: ignore
 
     xgboost_installed = True
 except ImportError:
     xgboost_installed = False
+    if TYPE_CHECKING:
+        import xgboost  # type: ignore
 
 
 def _rename_one_feature(name):