(torchx/entrypoints) Use importlib.metadata instead of the backported importlib_metadata

Differential Revision: D77619564

Pull Request resolved: meta-pytorch#1087

Author: kiukchung

requirements.txt

@@ -1,5 +1,4 @@
 docstring-parser>=0.8.1
-importlib-metadata
 pyyaml
 docker
 filelock

torchx/util/entrypoints.py

@@ -5,14 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
+# pyre-ignore-all-errors[3, 2, 16]
 
+from importlib import metadata
+from importlib.metadata import EntryPoint
 from typing import Any, Dict, Optional
 
-import importlib_metadata as metadata
-from importlib_metadata import EntryPoint
 
-
-# pyre-ignore-all-errors[3, 2]
 def load(group: str, name: str, default=None):
     """
     Loads the entry point specified by
@@ -30,13 +29,34 @@ def load(group: str, name: str, default=None):
     raises an error.
     """
 
-    entrypoints = metadata.entry_points().select(group=group)
+    # [note_on_entrypoints]
+    # return type of importlib.metadata.entry_points() is different between python-3.9 and python-3.10
+    # https://docs.python.org/3.9/library/importlib.metadata.html#importlib.metadata.entry_points
+    # https://docs.python.org/3.10/library/importlib.metadata.html#importlib.metadata.entry_points
+    if hasattr(metadata.entry_points(), "select"):
+        # python>=3.10
+        entrypoints = metadata.entry_points().select(group=group)
 
-    if name not in entrypoints.names and default is not None:
-        return default
+        if name not in entrypoints.names and default is not None:
+            return default
+
+        ep = entrypoints[name]
+        return ep.load()
 
-    ep = entrypoints[name]
-    return ep.load()
+    else:
+        # python<3.10 (e.g. 3.9)
+        # metadata.entry_points() returns dict[str, tuple[EntryPoint]] (not EntryPoints) in python-3.9
+        entrypoints = metadata.entry_points().get(group, ())
+
+        for ep in entrypoints:
+            if ep.name == name:
+                return ep.load()
+
+        # [group].name not found
+        if default is not None:
+            return default
+        else:
+            raise KeyError(f"entrypoint {group}.{name} not found")
 
 
 def _defer_load_ep(ep: EntryPoint) -> object:
@@ -49,7 +69,6 @@ def run(*args: object, **kwargs: object) -> object:
     return run
 
 
-# pyre-ignore-all-errors[3, 2]
 def load_group(
     group: str, default: Optional[Dict[str, Any]] = None, skip_defaults: bool = False
 ):
@@ -87,7 +106,13 @@ def load_group(
 
     """
 
-    entrypoints = metadata.entry_points().select(group=group)
+    # see [note_on_entrypoints] above
+    if hasattr(metadata.entry_points(), "select"):
+        # python>=3.10
+        entrypoints = metadata.entry_points().select(group=group)
+    else:
+        # python<3.10 (e.g. 3.9)
+        entrypoints = metadata.entry_points().get(group, ())
 
     if len(entrypoints) == 0:
         if skip_defaults:

torchx/util/test/entrypoints_test.py

@@ -8,16 +8,16 @@
 
 import unittest
 from configparser import ConfigParser
+
+from importlib.metadata import EntryPoint
 from types import ModuleType
-from typing import List
-from unittest.mock import MagicMock, patch
 
-from importlib_metadata import EntryPoint, EntryPoints
+from unittest.mock import MagicMock, patch
 
 from torchx.util.entrypoints import load, load_group
 
 
-def EntryPoint_from_config(config: ConfigParser) -> List[EntryPoint]:
+def EntryPoint_from_config(config: ConfigParser) -> list[EntryPoint]:
     # from stdlib, Copyright (c) Python Authors
     return [
         EntryPoint(name, value, group)
@@ -26,7 +26,7 @@ def EntryPoint_from_config(config: ConfigParser) -> List[EntryPoint]:
     ]
 
 
-def EntryPoint_from_text(text: str) -> List[EntryPoint]:
+def EntryPoint_from_text(text: str) -> list[EntryPoint]:
     # from stdlib, Copyright (c) Python Authors
     config = ConfigParser(delimiters="=")
     config.read_string(text)
@@ -66,14 +66,26 @@ def barbaz() -> str:
 [ep.grp.missing.mod.test]
 baz = torchx.util.test.entrypoints_test.missing_module
 """
-_ENTRY_POINTS: EntryPoints = EntryPoints(
+
+_EPS: list[EntryPoint] = (
     EntryPoint_from_text(_EP_TXT)
     + EntryPoint_from_text(_EP_GRP_TXT)
     + EntryPoint_from_text(_EP_GRP_IGN_ATTR_TXT)
     + EntryPoint_from_text(_EP_GRP_MOD_TXT)
     + EntryPoint_from_text(_EP_GRP_IGN_MOD_TXT)
 )
 
+try:
+    from importlib.metadata import EntryPoints
+except ImportError:
+    # python<=3.9
+    _ENTRY_POINTS: dict[str, list[EntryPoint]] = {}
+    for ep in _EPS:
+        _ENTRY_POINTS.setdefault(ep.group, []).append(ep)
+else:
+    # python>=3.10
+    _ENTRY_POINTS: EntryPoints = EntryPoints(_EPS)
+
 _METADATA_EPS: str = "torchx.util.entrypoints.metadata.entry_points"