Decrease start-up time of editable-installed entry points on newer versions of Python

3 files changed, 62 insertions, 26 deletions

diff --git a/changelog.d/2194.change.rst b/changelog.d/2194.change.rst
new file mode 100644
index 00000000..8b5adb31
--- /dev/null
+++ b/changelog.d/2194.change.rst
@@ -0,0 +1 @@
+Editable-installed entry points now load significantly faster on Python versions 3.8+.
diff --git a/setuptools/command/easy_install.py b/setuptools/command/easy_install.py
index 5a9576ff..a82b1655 100644
--- a/setuptools/command/easy_install.py
+++ b/setuptools/command/easy_install.py
@@ -2070,19 +2070,36 @@ class ScriptWriter:
     gui apps.
     """
 
-    template = textwrap.dedent(r"""
-        # EASY-INSTALL-ENTRY-SCRIPT: %(spec)r,%(group)r,%(name)r
-        __requires__ = %(spec)r
-        import re
-        import sys
-        from pkg_resources import load_entry_point
-
-        if __name__ == '__main__':
-            sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
-            sys.exit(
-                load_entry_point(%(spec)r, %(group)r, %(name)r)()
-            )
-    """).lstrip()
+    try:
+        from importlib.metadata import distribution  # noqa: F401
+
+        template = textwrap.dedent(r"""
+            # EASY-INSTALL-ENTRY-SCRIPT: %(spec)r,%(group)r,%(name)r
+            __requires__ = %(spec)r
+            import re
+            import sys
+            from importlib.metadata import distribution
+
+            if __name__ == '__main__':
+                sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
+                for entry_point in distribution(%(spec)r).entry_points:
+                    if entry_point.group == %(group)r and entry_point.name == %(name)r:
+                        sys.exit(entry_point.load()())
+        """).lstrip()  # noqa: E501
+    except ImportError:
+        template = textwrap.dedent(r"""
+            # EASY-INSTALL-ENTRY-SCRIPT: %(spec)r,%(group)r,%(name)r
+            __requires__ = %(spec)r
+            import re
+            import sys
+            from pkg_resources import load_entry_point
+
+            if __name__ == '__main__':
+                sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
+                sys.exit(
+                    load_entry_point(%(spec)r, %(group)r, %(name)r)()
+                )
+        """).lstrip()  # noqa: E501
 
     command_spec_class = CommandSpec
 
diff --git a/setuptools/tests/test_easy_install.py b/setuptools/tests/test_easy_install.py
index 3044cbd0..58108db6 100644
--- a/setuptools/tests/test_easy_install.py
+++ b/setuptools/tests/test_easy_install.py
@@ -71,19 +71,37 @@ class TestEasyInstallTest:
 
     def test_get_script_args(self):
         header = ei.CommandSpec.best().from_environment().as_header()
-        expected = header + DALS(r"""
-            # EASY-INSTALL-ENTRY-SCRIPT: 'spec','console_scripts','name'
-            __requires__ = 'spec'
-            import re
-            import sys
-            from pkg_resources import load_entry_point
-
-            if __name__ == '__main__':
-                sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
-                sys.exit(
-                    load_entry_point('spec', 'console_scripts', 'name')()
-                )
-            """)  # noqa: E501
+        try:
+            from importlib.metadata import distribution  # noqa: F401
+
+            expected = header + DALS(r"""
+                # EASY-INSTALL-ENTRY-SCRIPT: 'spec','console_scripts','name'
+                __requires__ = 'spec'
+                import re
+                import sys
+                from importlib.metadata import distribution
+
+                if __name__ == '__main__':
+                    sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
+                    for entry_point in distribution('spec').entry_points:
+                        if entry_point.group == 'console_scripts' and entry_point.name == 'name':
+                            sys.exit(entry_point.load()())
+                """)  # noqa: E501
+        except ImportError:
+            expected = header + DALS(r"""
+                # EASY-INSTALL-ENTRY-SCRIPT: 'spec','console_scripts','name'
+                __requires__ = 'spec'
+                import re
+                import sys
+                from pkg_resources import load_entry_point
+
+                if __name__ == '__main__':
+                    sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
+                    sys.exit(
+                        load_entry_point('spec', 'console_scripts', 'name')()
+                    )
+                """)  # noqa: E501
+
         dist = FakeDist()
 
         args = next(ei.ScriptWriter.get_args(dist))