[SETUP] Add optional dependencies to extras_require (#4428)

diff --git a/python/setup.py b/python/setup.py
index ad14df1..bc3390a 100644 (file)
--- a/python/setup.py
+++ b/python/setup.py
@@ -36,6 +36,7 @@ else:
 
 CURRENT_DIR = os.path.dirname(__file__)
 
+
 def get_lib_path():
     """Get library path, name and version"""
     # We can not import `libinfo.py` in setup.py directly since __init__.py
@@ -56,8 +57,10 @@ def get_lib_path():
         libs = None
     return libs, version
 
+
 LIB_LIST, __version__ = get_lib_path()
 
+
 def config_cython():
     """Try to configure cython and return cython configuration"""
     if os.name == 'nt':
@@ -101,6 +104,7 @@ def config_cython():
         print("WARNING: Cython is not installed, will compile without cython module")
         return []
 
+
 class BinaryDistribution(Distribution):
     def has_ext_modules(self):
         return True
@@ -108,6 +112,7 @@ class BinaryDistribution(Distribution):
     def is_pure(self):
         return False
 
+
 include_libs = False
 wheel_include_libs = False
 if not os.getenv('CONDA_BUILD'):
@@ -138,10 +143,12 @@ if include_libs:
         "data_files": [('tvm', LIB_LIST)]
     }
 
+
 def get_package_data_files():
     # Relay standard libraries
     return ['relay/std/prelude.rly']
 
+
 setup(name='tvm',
       version=__version__,
       description="TVM: An End to End Tensor IR/DSL Stack for Deep Learning Systems",
@@ -152,6 +159,15 @@ setup(name='tvm',
         'attrs',
         'psutil',
         ],
+      extras_require={'test': ['PIL',
+                               'matplotlib'],
+                      'extra_feature': ['tornado',
+                                        'psutil',
+                                        'xgboost',
+                                        'mypy',
+                                        'orderedset',
+                                        'antlr4-python3-runtime']},
+
       packages=find_packages(),
       package_dir={'tvm': 'tvm'},
       package_data={'tvm': get_package_data_files()},