_loader.py 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. from __future__ import annotations
  2. import importlib.metadata as importlib_metadata
  3. import os
  4. import warnings
  5. from collections.abc import Iterable
  6. from typing import TYPE_CHECKING, Final
  7. if TYPE_CHECKING:
  8. from . import PydanticPluginProtocol
  9. PYDANTIC_ENTRY_POINT_GROUP: Final[str] = 'pydantic'
  10. # cache of plugins
  11. _plugins: dict[str, PydanticPluginProtocol] | None = None
  12. # return no plugins while loading plugins to avoid recursion and errors while import plugins
  13. # this means that if plugins use pydantic
  14. _loading_plugins: bool = False
  15. def get_plugins() -> Iterable[PydanticPluginProtocol]:
  16. """Load plugins for Pydantic.
  17. Inspired by: https://github.com/pytest-dev/pluggy/blob/1.3.0/src/pluggy/_manager.py#L376-L402
  18. """
  19. disabled_plugins = os.getenv('PYDANTIC_DISABLE_PLUGINS')
  20. global _plugins, _loading_plugins
  21. if _loading_plugins:
  22. # this happens when plugins themselves use pydantic, we return no plugins
  23. return ()
  24. elif disabled_plugins in ('__all__', '1', 'true'):
  25. return ()
  26. elif _plugins is None:
  27. _plugins = {}
  28. # set _loading_plugins so any plugins that use pydantic don't themselves use plugins
  29. _loading_plugins = True
  30. try:
  31. for dist in importlib_metadata.distributions():
  32. for entry_point in dist.entry_points:
  33. if entry_point.group != PYDANTIC_ENTRY_POINT_GROUP:
  34. continue
  35. if entry_point.value in _plugins:
  36. continue
  37. if disabled_plugins is not None and entry_point.name in disabled_plugins.split(','):
  38. continue
  39. try:
  40. _plugins[entry_point.value] = entry_point.load()
  41. except (ImportError, AttributeError) as e:
  42. warnings.warn(
  43. f'{e.__class__.__name__} while loading the `{entry_point.name}` Pydantic plugin, '
  44. f'this plugin will not be installed.\n\n{e!r}'
  45. )
  46. finally:
  47. _loading_plugins = False
  48. return _plugins.values()