_loader.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  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. stacklevel=2,
  46. )
  47. finally:
  48. _loading_plugins = False
  49. return _plugins.values()