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