call_context.py 3.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. """Provides access to variables pertaining to specific call contexts."""
  2. # Copyright (c) Jupyter Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. from contextvars import Context, ContextVar, copy_context
  5. from typing import Any
  6. class CallContext:
  7. """CallContext essentially acts as a namespace for managing context variables.
  8. Although not required, it is recommended that any "file-spanning" context variable
  9. names (i.e., variables that will be set or retrieved from multiple files or services) be
  10. added as constants to this class definition.
  11. """
  12. # Add well-known (file-spanning) names here.
  13. #: Provides access to the current request handler once set.
  14. JUPYTER_HANDLER: str = "JUPYTER_HANDLER"
  15. # A map of variable name to value is maintained as the single ContextVar. This also enables
  16. # easier management over maintaining a set of ContextVar instances, since the Context is a
  17. # map of ContextVar instances to their values, and the "name" is no longer a lookup key.
  18. _NAME_VALUE_MAP = "_name_value_map"
  19. _name_value_map: ContextVar[dict[str, Any]] = ContextVar(_NAME_VALUE_MAP)
  20. @classmethod
  21. def get(cls, name: str) -> Any:
  22. """Returns the value corresponding the named variable relative to this context.
  23. If the named variable doesn't exist, None will be returned.
  24. Parameters
  25. ----------
  26. name : str
  27. The name of the variable to get from the call context
  28. Returns
  29. -------
  30. value: Any
  31. The value associated with the named variable for this call context
  32. """
  33. name_value_map = CallContext._get_map()
  34. if name in name_value_map:
  35. return name_value_map[name]
  36. return None # TODO: should this raise `LookupError` (or a custom error derived from said)
  37. @classmethod
  38. def set(cls, name: str, value: Any) -> None:
  39. """Sets the named variable to the specified value in the current call context.
  40. Parameters
  41. ----------
  42. name : str
  43. The name of the variable to store into the call context
  44. value : Any
  45. The value of the variable to store into the call context
  46. Returns
  47. -------
  48. None
  49. """
  50. name_value_map = CallContext._get_map()
  51. name_value_map[name] = value
  52. @classmethod
  53. def context_variable_names(cls) -> list[str]:
  54. """Returns a list of variable names set for this call context.
  55. Returns
  56. -------
  57. names: List[str]
  58. A list of variable names set for this call context.
  59. """
  60. name_value_map = CallContext._get_map()
  61. return list(name_value_map.keys())
  62. @classmethod
  63. def _get_map(cls) -> dict[str, Any]:
  64. """Get the map of names to their values from the _NAME_VALUE_MAP context var.
  65. If the map does not exist in the current context, an empty map is created and returned.
  66. """
  67. ctx: Context = copy_context()
  68. if CallContext._name_value_map not in ctx:
  69. CallContext._name_value_map.set({})
  70. return CallContext._name_value_map.get()