Improve startup time with lazy loading some args (#1221)

* Improve startup time with lazy loading some args

* add some tests

* Add changelog entry

* Update CHANGELOG.md

Co-authored-by: Jakub Roztocil <jakub@roztocil.co>

CHANGELOG.md

@@ -9,6 +9,7 @@ This project adheres to [Semantic Versioning](https://semver.org/).
 - Added support for receving multiple HTTP headers with the same name, individually. ([#1207](https://github.com/httpie/httpie/issues/1207))
 - Added support for keeping `://` in the URL argument to allow quick conversions of pasted URLs into HTTPie calls just by adding a space after the protocol name (`$ https ://pie.dev` → `https://pie.dev`). ([#1195](https://github.com/httpie/httpie/issues/1195))
 - Added support for basic JSON types on `--form`/`--multipart` when using JSON only operators (`:=`/`:=@`). ([#1212](https://github.com/httpie/httpie/issues/1212))
+- Improved startup time by 40% with lazily loading pygments plugins. ([#1211](https://github.com/httpie/httpie/pull/1211))
 - Added support for `bearer` authentication method ([#1215](https://github.com/httpie/httpie/issues/1215)).
 
 ## [2.6.0](https://github.com/httpie/httpie/compare/2.5.0...2.6.0) (2021-10-14)

httpie/cli/definition.py

@@ -19,8 +19,9 @@ from .constants import (
     SORTED_FORMAT_OPTIONS_STRING,
     UNSORTED_FORMAT_OPTIONS_STRING,
 )
+from .utils import LazyChoices
 from ..output.formatters.colors import (
-    AUTO_STYLE, AVAILABLE_STYLES, DEFAULT_STYLE,
+    AUTO_STYLE, DEFAULT_STYLE, get_available_styles
 )
 from ..plugins.builtin import BuiltinAuthPlugin
 from ..plugins.registry import plugin_manager
@@ -41,6 +42,7 @@ parser = HTTPieArgumentParser(
 
     '''),
 )
+parser.register('action', 'lazy_choices', LazyChoices)
 
 #######################################################################
 # Positional arguments.
@@ -247,32 +249,38 @@ output_processing.add_argument(
 
     '''
 )
+
+
+def format_style_help(available_styles):
+    return '''
+    Output coloring style (default is "{default}"). It can be one of:
+
+        {available_styles}
+
+    The "{auto_style}" style follows your terminal's ANSI color styles.
+    For non-{auto_style} styles to work properly, please make sure that the
+    $TERM environment variable is set to "xterm-256color" or similar
+    (e.g., via `export TERM=xterm-256color' in your ~/.bashrc).
+    '''.format(
+        default=DEFAULT_STYLE,
+        available_styles='\n'.join(
+            f'        {line.strip()}'
+            for line in wrap(', '.join(available_styles), 60)
+        ).strip(),
+        auto_style=AUTO_STYLE,
+    )
+
+
 output_processing.add_argument(
     '--style', '-s',
     dest='style',
     metavar='STYLE',
     default=DEFAULT_STYLE,
-    choices=sorted(AVAILABLE_STYLES),
-    help='''
-    Output coloring style (default is "{default}"). It can be One of:
-
-        {available_styles}
-
-    The "{auto_style}" style follows your terminal's ANSI color styles.
-
-    For non-{auto_style} styles to work properly, please make sure that the
-    $TERM environment variable is set to "xterm-256color" or similar
-    (e.g., via `export TERM=xterm-256color' in your ~/.bashrc).
-
-    '''.format(
-        default=DEFAULT_STYLE,
-        available_styles='\n'.join(
-            f'        {line.strip()}'
-            for line in wrap(', '.join(sorted(AVAILABLE_STYLES)), 60)
-        ).strip(),
-        auto_style=AUTO_STYLE,
-    )
+    action='lazy_choices',
+    getter=get_available_styles,
+    help_formatter=format_style_help
 )
+
 _sorted_kwargs = {
     'action': 'append_const',
     'const': SORTED_FORMAT_OPTIONS_STRING,
@@ -564,27 +572,14 @@ auth.add_argument(
 )
 
 
-class _AuthTypeLazyChoices:
-    # Needed for plugin testing
-
-    def __contains__(self, item):
-        return item in plugin_manager.get_auth_plugin_mapping()
-
-    def __iter__(self):
-        return iter(sorted(plugin_manager.get_auth_plugin_mapping().keys()))
-
-
-_auth_plugins = plugin_manager.get_auth_plugins()
-auth.add_argument(
-    '--auth-type', '-A',
-    choices=_AuthTypeLazyChoices(),
-    default=None,
-    help='''
+def format_auth_help(auth_plugins_mapping):
+    auth_plugins = list(auth_plugins_mapping.values())
+    return '''
     The authentication mechanism to be used. Defaults to "{default}".
 
     {types}
 
-    '''.format(default=_auth_plugins[0].auth_type, types='\n    '.join(
+    '''.format(default=auth_plugins[0].auth_type, types='\n    '.join(
         '"{type}": {name}{package}{description}'.format(
             type=plugin.auth_type,
             name=plugin.name,
@@ -597,8 +592,18 @@ auth.add_argument(
                 '\n      ' + ('\n      '.join(wrap(plugin.description)))
             )
         )
-        for plugin in _auth_plugins
-    )),
+        for plugin in auth_plugins
+    ))
+
+
+auth.add_argument(
+    '--auth-type', '-A',
+    action='lazy_choices',
+    default=None,
+    getter=plugin_manager.get_auth_plugin_mapping,
+    sort=True,
+    cache=False,
+    help_formatter=format_auth_help,
 )
 auth.add_argument(
     '--ignore-netrc',

httpie/cli/utils.py

@@ -0,0 +1,53 @@
+import argparse
+from typing import Any, Callable, Generic, Iterator, Iterable, Optional, TypeVar
+
+T = TypeVar('T')
+
+
+class LazyChoices(argparse.Action, Generic[T]):
+    def __init__(
+        self,
+        *args,
+        getter: Callable[[], Iterable[T]],
+        help_formatter: Optional[Callable[[T], str]] = None,
+        sort: bool = False,
+        cache: bool = True,
+        **kwargs
+    ) -> None:
+        self.getter = getter
+        self.help_formatter = help_formatter
+        self.sort = sort
+        self.cache = cache
+        self._help: Optional[str] = None
+        self._obj: Optional[Iterable[T]] = None
+        super().__init__(*args, **kwargs)
+        self.choices = self
+
+    def load(self) -> T:
+        if self._obj is None or not self.cache:
+            self._obj = self.getter()
+
+        assert self._obj is not None
+        return self._obj
+
+    @property
+    def help(self) -> str:
+        if self._help is None and self.help_formatter is not None:
+            self._help = self.help_formatter(self.load())
+        return self._help
+
+    @help.setter
+    def help(self, value: Any) -> None:
+        self._help = value
+
+    def __contains__(self, item: Any) -> bool:
+        return item in self.load()
+
+    def __iter__(self) -> Iterator[T]:
+        if self.sort:
+            return iter(sorted(self.load()))
+        else:
+            return iter(self.load())
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, values)

httpie/output/formatters/colors.py

@@ -28,9 +28,14 @@ if is_windows:
     # great and fruity seems to give the best result there.
     DEFAULT_STYLE = 'fruity'
 
-AVAILABLE_STYLES = set(pygments.styles.get_all_styles())
-AVAILABLE_STYLES.add(SOLARIZED_STYLE)
-AVAILABLE_STYLES.add(AUTO_STYLE)
+BUNDLED_STYLES = {
+    SOLARIZED_STYLE,
+    AUTO_STYLE
+}
+
+
+def get_available_styles():
+    return BUNDLED_STYLES | set(pygments.styles.get_all_styles())
 
 
 class ColorFormatter(FormatterPlugin):

tests/test_cli_utils.py

@@ -0,0 +1,86 @@
+import pytest
+from argparse import ArgumentParser
+from unittest.mock import Mock
+from httpie.cli.utils import LazyChoices
+
+
+def test_lazy_choices():
+    mock = Mock()
+    getter = mock.getter
+    getter.return_value = ['a', 'b', 'c']
+
+    parser = ArgumentParser()
+    parser.register('action', 'lazy_choices', LazyChoices)
+    parser.add_argument(
+        '--option',
+        help="the regular option",
+        default='a',
+        metavar='SYMBOL',
+        choices=['a', 'b'],
+    )
+    parser.add_argument(
+        '--lazy-option',
+        help="the lazy option",
+        default='a',
+        metavar='SYMBOL',
+        action='lazy_choices',
+        getter=getter,
+        cache=False  # for test purposes
+    )
+
+    # Parser initalization doesn't call it.
+    getter.assert_not_called()
+
+    # If we don't use --lazy-option, we don't retrieve it.
+    parser.parse_args([])
+    getter.assert_not_called()
+
+    parser.parse_args(['--option', 'b'])
+    getter.assert_not_called()
+
+    # If we pass a value, it will retrieve to verify.
+    parser.parse_args(['--lazy-option', 'c'])
+    getter.assert_called()
+    getter.reset_mock()
+
+    with pytest.raises(SystemExit):
+        parser.parse_args(['--lazy-option', 'z'])
+    getter.assert_called()
+    getter.reset_mock()
+
+
+def test_lazy_choices_help():
+    mock = Mock()
+    getter = mock.getter
+    getter.return_value = ['a', 'b', 'c']
+
+    help_formatter = mock.help_formatter
+    help_formatter.return_value = '<my help>'
+
+    parser = ArgumentParser()
+    parser.register('action', 'lazy_choices', LazyChoices)
+    parser.add_argument(
+        '--lazy-option',
+        default='a',
+        metavar='SYMBOL',
+        action='lazy_choices',
+        getter=getter,
+        help_formatter=help_formatter,
+        cache=False  # for test purposes
+    )
+
+    # Parser initalization doesn't call it.
+    getter.assert_not_called()
+
+    # If we don't use `--help`, we don't use it.
+    parser.parse_args([])
+    getter.assert_not_called()
+    help_formatter.assert_not_called()
+
+    parser.parse_args(['--lazy-option', 'b'])
+    help_formatter.assert_not_called()
+
+    # If we use --help, then we call it with styles
+    with pytest.raises(SystemExit):
+        parser.parse_args(['--help'])
+    help_formatter.assert_called_once_with(['a', 'b', 'c'])