From 9f093c3f1f17721e7195aae99d57618e98ae50a6 Mon Sep 17 00:00:00 2001
From: Keith Battocchi <kebatt@microsoft.com>
Date: Thu, 18 Feb 2021 21:46:57 -0500
Subject: [PATCH] Create new extras

---
 azure-pipelines.yml                       |  4 +++-
 econml/cate_interpreter/_interpreters.py  |  8 +++++++-
 econml/cate_interpreter/_tree_exporter.py | 12 +++++++++---
 econml/iv/nnet/_deepiv.py                 | 14 ++++++++++----
 setup.cfg                                 | 16 ++++++++++++----
 5 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 8364be9b8..96a5bd720 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -78,7 +78,7 @@ jobs:
 
         - script: 'python setup.py build_sphinx -b doctest'
           displayName: 'Run doctests'
-      package: '-e .[automl]'
+      package: '-e .[all]'
 
 - job: 'Notebooks'
   dependsOn: 'EvalChanges'
@@ -106,6 +106,7 @@ jobs:
             testResultsFiles: '**/test-results.xml'
             testRunTitle: 'Notebooks'
           condition: succeededOrFailed()
+      package: '-e .[tf,interpret]'
 
 # - job: 'AutoML'
 #   dependsOn: 'EvalChanges'
@@ -216,3 +217,4 @@ jobs:
           inputs:
             codeCoverageTool: Cobertura
             summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'
+      package: '-e .[tf,interpret]'
diff --git a/econml/cate_interpreter/_interpreters.py b/econml/cate_interpreter/_interpreters.py
index f3147f373..5fbed5045 100644
--- a/econml/cate_interpreter/_interpreters.py
+++ b/econml/cate_interpreter/_interpreters.py
@@ -6,9 +6,15 @@
 from io import StringIO
 from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier
 from sklearn.utils.validation import check_is_fitted
-import graphviz
 from ._tree_exporter import _CateTreeDOTExporter, _CateTreeMPLExporter, _PolicyTreeDOTExporter, _PolicyTreeMPLExporter
 
+try:
+    import graphviz
+except ImportError as exn:
+    raise ImportError("matplotlib and graphviz are no longer dependencies of the main econml package;\n"
+                      "install econml[interpret] or econml[all] to require them, or install them separately, "
+                      "to use the tree interpreters") from exn
+
 
 class _SingleTreeInterpreter(metaclass=abc.ABCMeta):
 
diff --git a/econml/cate_interpreter/_tree_exporter.py b/econml/cate_interpreter/_tree_exporter.py
index b4232493e..3c1192e7d 100644
--- a/econml/cate_interpreter/_tree_exporter.py
+++ b/econml/cate_interpreter/_tree_exporter.py
@@ -3,9 +3,15 @@
 
 import numpy as np
 import re
-import matplotlib
-matplotlib.use('Agg')
-import matplotlib.pyplot as plt
+try:
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+except ImportError as exn:
+    raise ImportError("matplotlib and graphviz are no longer dependencies of the main econml package;\n"
+                      "install econml[interpret] or econml[all] to require them, or install them separately, "
+                      "to use the tree interpreters") from exn
+
 
 # HACK: We're relying on some of sklearn's non-public classes which are not completely stable.
 #       However, the alternative is reimplementing a bunch of intricate stuff by hand
diff --git a/econml/iv/nnet/_deepiv.py b/econml/iv/nnet/_deepiv.py
index 6f5a8369b..8622d8945 100644
--- a/econml/iv/nnet/_deepiv.py
+++ b/econml/iv/nnet/_deepiv.py
@@ -4,13 +4,19 @@
 """Deep IV estimator and related components."""
 
 import numpy as np
-import keras
 from ..._cate_estimator import BaseCateEstimator
 from ...utilities import deprecated
-from keras import backend as K
-import keras.layers as L
-from keras.models import Model
 from econml.utilities import check_input_arrays, _deprecate_positional
+try:
+    import keras
+    from keras import backend as K
+    import keras.layers as L
+    from keras.models import Model
+except ImportError as exn:
+    raise ImportError("keras and tensorflow are no longer dependencies of the main econml package;\n"
+                      "install econml[tf] or econml[all] to require them, or install them separately, "
+                      "to use DeepIv") from exn
+
 
 # TODO: make sure to use random seeds wherever necessary
 # TODO: make sure that the public API consistently uses "T" instead of "P" for the treatment
diff --git a/setup.cfg b/setup.cfg
index c7bfd48d4..b916e71b7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -43,14 +43,10 @@ install_requires =
     numpy
     scipy > 1.4.0
     scikit-learn >= 0.24
-    keras < 2.4
     sparse
-    tensorflow > 1.10, < 2.3
     joblib >= 0.13.0
     numba != 0.42.1
     statsmodels >= 0.9
-    graphviz
-    matplotlib
     pandas
     shap ~= 0.38.1
     dowhy
@@ -71,6 +67,18 @@ automl =
     ; Disabled due to incompatibility with scikit-learn
     ; azureml-sdk[explain,automl] == 1.0.83
     azure-cli
+tf =
+    keras < 2.4
+    tensorflow > 1.10, < 2.3
+interpret =
+    matplotlib
+    graphviz
+all =
+    azure-cli
+    keras < 2.4
+    tensorflow > 1.10, < 2.3
+    matplotlib
+    graphviz
     
 ; TODO: exclude tests?
 [options.packages.find]