(torchx/components) Remove dependency to pyre_extensions (#1088)

Summary:

There was only a single usage of `pyre_extensions` in the entire code base. Get rid of this dependency by using an `assert` statement rather than `pyre_extensions.none_throws`.

Reviewed By: highker

Differential Revision: D77621041

Author: kiukchung

requirements.txt

@@ -1,4 +1,3 @@
-pyre-extensions
 docstring-parser>=0.8.1
 importlib-metadata
 pyyaml

torchx/components/structured_arg.py

@@ -30,8 +30,6 @@
 from pathlib import Path
 from typing import Optional
 
-from pyre_extensions import none_throws
-
 from torchx import specs
 
 
@@ -148,7 +146,8 @@ def parse_from(
             if m:  # use the last module name
                 run_name = m.rpartition(".")[2]
             else:  # use script name w/ no extension
-                run_name = Path(none_throws(script)).stem
+                assert script, "`script` can't be `None` here due checks above"
+                run_name = Path(script).stem
         return StructuredNameArgument(
             experiment_name or default_experiment_name, run_name
         )

torchx/util/test/types_test.py

@@ -8,9 +8,8 @@
 
 import inspect
 import unittest
-from typing import cast, Dict, List, Optional, Union
+from typing import cast, Optional, Union
 
-import typing_inspect
 from torchx.util.types import (
     decode,
     decode_from_string,
@@ -26,25 +25,25 @@
 
 def _test_complex_args(
     arg1: int,
-    arg2: Optional[List[str]],
+    arg2: Optional[list[str]],
     arg3: Union[float, int],
 ) -> int:
     return 42
 
 
-def _test_dict(arg1: Dict[int, float]) -> int:
+def _test_dict(arg1: dict[int, float]) -> int:
     return 42
 
 
-def _test_nested_object(arg1: Dict[str, List[str]]) -> int:
+def _test_nested_object(arg1: dict[str, list[str]]) -> int:
     return 42
 
 
-def _test_list(arg1: List[float]) -> int:
+def _test_list(arg1: list[float]) -> int:
     return 42
 
 
-def _test_complex_list(arg1: List[List[float]]) -> int:
+def _test_complex_list(arg1: list[list[float]]) -> int:
     return 42
 
 
@@ -59,24 +58,21 @@ def test_decode_optional(self) -> None:
         arg1_parameter = parameters["arg1"]
         arg1_type = decode_optional(arg1_parameter.annotation)
         self.assertTrue(arg1_type is int)
-
-        arg2_parameter = parameters["arg2"]
         arg2_type = decode_optional(parameters["arg2"].annotation)
-        self.assertTrue(typing_inspect.get_origin(arg2_type) is list)
-
+        self.assertTrue(getattr(arg2_type, "__origin__", None) is list)
         arg3_parameter = parameters["arg3"]
         arg3_type = decode_optional(arg3_parameter.annotation)
-        self.assertTrue(typing_inspect.get_origin(arg3_type) is Union)
+        self.assertTrue(
+            hasattr(arg3_type, "__origin__") and arg3_type.__origin__ is Union
+        )
 
     def test_is_primitive(self) -> None:
         parameters = inspect.signature(_test_complex_args).parameters
 
         arg1_parameter = parameters["arg1"]
-        arg1_type = decode_optional(arg1_parameter.annotation)
         self.assertTrue(is_primitive(arg1_parameter.annotation))
 
         arg2_parameter = parameters["arg2"]
-        arg2_type = decode_optional(parameters["arg2"].annotation)
         self.assertFalse(is_primitive(arg2_parameter.annotation))
 
     def test_is_bool(self) -> None:
@@ -89,7 +85,7 @@ def test_decode_from_string_dict(self) -> None:
         encoded_value = "1=1.0,2=42.1,3=10"
 
         value = decode_from_string(encoded_value, parameters["arg1"].annotation)
-        value = cast(Dict[int, float], value)
+        value = cast(dict[int, float], value)
         self.assertEqual(3, len(value))
         self.assertEqual(float(1.0), value[1])
         self.assertEqual(float(42.1), value[2])
@@ -101,7 +97,7 @@ def test_decode_from_string_list(self) -> None:
         encoded_value = "1.0,42.2,3.9"
 
         value = decode_from_string(encoded_value, parameters["arg1"].annotation)
-        value = cast(List[float], value)
+        value = cast(list[float], value)
         self.assertEqual(3, len(value))
         self.assertEqual(float(1.0), value[0])
         self.assertEqual(float(42.2), value[1])
@@ -217,8 +213,8 @@ def fake_component(
             f: float,
             s: str,
             b: bool,
-            l: List[str],
-            m: Dict[str, str],
+            l: list[str],
+            m: dict[str, str],
             o: Optional[int],
         ) -> None:
             # component has to return an AppDef

torchx/util/types.py

@@ -8,12 +8,10 @@
 
 import inspect
 import re
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Callable, Optional, Tuple, TypeVar, Union
 
-import typing_inspect
 
-
-def to_list(arg: str) -> List[str]:
+def to_list(arg: str) -> list[str]:
     conf = []
     if len(arg.strip()) == 0:
         return []
@@ -22,9 +20,9 @@ def to_list(arg: str) -> List[str]:
     return conf
 
 
-def to_dict(arg: str) -> Dict[str, str]:
+def to_dict(arg: str) -> dict[str, str]:
     """
-    Parses the given ``arg`` string literal into a ``Dict[str, str]`` of
+    Parses the given ``arg`` string literal into a ``dict[str, str]`` of
     key-value pairs delimited by ``"="`` (equals). The values may be a
     list literal where the list elements are delimited by ``","`` (comma)
     or ``";"`` (semi-colon). The same delimiters (``","`` and ``";"``) are used
@@ -85,14 +83,14 @@ def to_val(val: str) -> str:
             return val[1:-1]
         return val if val != '""' and val != "''" else ""
 
-    arg_map: Dict[str, str] = {}
+    arg_map: dict[str, str] = {}
 
     if not arg:
         return arg_map
 
     # find quoted values
     quoted_pattern = r'([\'"])((?:\\.|(?!\1).)*?)\1'
-    quoted_values: List[str] = []
+    quoted_values: list[str] = []
 
     def replace_quoted(match):
         quoted_values.append(match.group(0))
@@ -133,19 +131,26 @@ def replace_quoted(match):
 
 # pyre-ignore-all-errors[3, 2]
 def _decode_string_to_dict(
-    encoded_value: str, param_type: Type[Dict[Any, Any]]
-) -> Dict[Any, Any]:
-    key_type, value_type = typing_inspect.get_args(param_type)
+    encoded_value: str, param_type: type[dict[Any, Any]]
+) -> dict[Any, Any]:
+    # pyre-ignore[16]
+    if not hasattr(param_type, "__args__") or len(param_type.__args__) != 2:
+        raise ValueError(f"param_type must be a `dict` type, but was `{param_type}`")
+
+    key_type, value_type = param_type.__args__
     arg_values = {}
     for key, value in to_dict(encoded_value).items():
         arg_values[key_type(key)] = value_type(value)
     return arg_values
 
 
 def _decode_string_to_list(
-    encoded_value: str, param_type: Type[List[Any]]
-) -> List[Any]:
-    value_type = typing_inspect.get_args(param_type)[0]
+    encoded_value: str, param_type: type[list[Any]]
+) -> list[Any]:
+    # pyre-ignore[16]
+    if not hasattr(param_type, "__args__") or len(param_type.__args__) != 1:
+        raise ValueError(f"param_type must be a `list` type, but was `{param_type}`")
+    value_type = param_type.__args__[0]
     if not is_primitive(value_type):
         raise ValueError("List types support only primitives: int, str, float")
     arg_values = []
@@ -166,7 +171,7 @@ def decode(encoded_value: Any, annotation: Any):
 
 def decode_from_string(
     encoded_value: str, annotation: Any
-) -> Union[Dict[Any, Any], List[Any], None]:
+) -> Union[dict[Any, Any], list[Any], None]:
     """Decodes string representation to the underlying type(Dict or List)
 
     Given a string representation of the value, the method decodes it according
@@ -191,13 +196,13 @@ def decode_from_string(
     if not encoded_value:
         return None
     value_type = annotation
-    value_origin = typing_inspect.get_origin(value_type)
-    if value_origin is dict:
-        return _decode_string_to_dict(encoded_value, value_type)
-    elif value_origin is list:
-        return _decode_string_to_list(encoded_value, value_type)
-    else:
-        raise ValueError("Unknown")
+    if hasattr(value_type, "__origin__"):
+        value_origin = value_type.__origin__
+        if value_origin is dict:
+            return _decode_string_to_dict(encoded_value, value_type)
+        elif value_origin is list:
+            return _decode_string_to_list(encoded_value, value_type)
+    raise ValueError("Unknown")
 
 
 def is_bool(param_type: Any) -> bool:
@@ -229,12 +234,13 @@ def decode_optional(param_type: Any) -> Any:
         If ``param_type`` is type Optional[INNER_TYPE], method returns INNER_TYPE
         Otherwise returns ``param_type``
     """
-    param_origin = typing_inspect.get_origin(param_type)
-    if param_origin is not Union:
+    if not hasattr(param_type, "__origin__"):
+        return param_type
+    if param_type.__origin__ is not Union:
         return param_type
-    key_type, value_type = typing_inspect.get_args(param_type)
-    if value_type is type(None):
-        return key_type
+    args = param_type.__args__
+    if len(args) == 2 and args[1] is type(None):
+        return args[0]
     else:
         return param_type