Replace nni.typehint with typing_extensions (#4802)

nni/algorithms/hpo/medianstop_assessor.py

@@ -4,11 +4,12 @@
 from __future__ import annotations
 
 import logging
+
 from schema import Schema, Optional
+from typing_extensions import Literal
 
 from nni import ClassArgsValidator
 from nni.assessor import Assessor, AssessResult
-from nni.typehint import Literal
 from nni.utils import extract_scalar_history
 
 logger = logging.getLogger('medianstop_Assessor')

nni/algorithms/hpo/tpe_tuner.py

@@ -22,10 +22,10 @@ from typing import Any, NamedTuple
 
 import numpy as np
 from scipy.special import erf  # pylint: disable=no-name-in-module
+from typing_extensions import Literal
 
 from nni.common.hpo_utils import Deduplicator, OptimizeMode, format_search_space, deformat_parameters, format_parameters
 from nni.tuner import Tuner
-from nni.typehint import Literal
 from nni.utils import extract_scalar_reward
 from . import random_tuner
 

nni/retiarii/strategy/_rl_impl.py

@@ -19,7 +19,7 @@ from gym import spaces
 from tianshou.data import to_torch
 from tianshou.env.worker import EnvWorker
 
-from nni.typehint import TypedDict
+from typing_extensions import TypedDict
 
 from .utils import get_targeted_model
 from ..graph import ModelStatus

nni/tools/package_utils/common.py

@@ -7,7 +7,7 @@ __all__ = ['AlgoMeta']
 
 from typing import NamedTuple
 
-from nni.typehint import Literal
+from typing_extensions import Literal
 
 class AlgoMeta(NamedTuple):
     name: str

nni/tools/package_utils/tuner_factory.py

@@ -13,7 +13,8 @@ import os
 import sys
 from typing import Any
 
-from nni.typehint import Literal
+from typing_extensions import Literal
+
 from . import config_manager
 
 ALGO_TYPES = ['tuners', 'assessors']

nni/typehint.py

@@ -5,18 +5,11 @@
 Types for static checking.
 """
 
-__all__ = [
-    'Literal', 'TypedDict',
-    'Parameters', 'SearchSpace', 'TrialMetric', 'TrialRecord',
-]
-
-import sys
-from typing import Any, Dict, List, TYPE_CHECKING
-
-if TYPE_CHECKING or sys.version_info >= (3, 8):
-    from typing import Literal, TypedDict
-else:
-    from typing_extensions import Literal, TypedDict
+__all__ = ['Parameters', 'SearchSpace', 'TrialMetric', 'TrialRecord']
+
+from typing import Any, Dict, List
+
+from typing_extensions import Literal, TypedDict
 
 Parameters = Dict[str, Any]
 """