Coverage for databooks/affirm.py: 98%
84 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 15:04 +0200
1"""Functions to safely evaluate strings and inspect notebook."""
3import ast
4from collections import abc
5from copy import deepcopy
6from itertools import compress
7from pathlib import Path
8from typing import Any, Callable, Dict, Iterable, List, Tuple
10from databooks import JupyterNotebook
11from databooks.data_models.base import DatabooksBase
12from databooks.logging import get_logger, set_verbose
14logger = get_logger(__file__)
16_ALLOWED_BUILTINS = (
17 all,
18 any,
19 enumerate,
20 filter,
21 getattr,
22 hasattr,
23 len,
24 list,
25 range,
26 sorted,
27)
28_ALLOWED_NODES = (
29 ast.Add,
30 ast.And,
31 ast.BinOp,
32 ast.BitAnd,
33 ast.BitOr,
34 ast.BitXor,
35 ast.BoolOp,
36 ast.boolop,
37 ast.cmpop,
38 ast.Compare,
39 ast.comprehension,
40 ast.Constant,
41 ast.Dict,
42 ast.Eq,
43 ast.Expr,
44 ast.expr,
45 ast.expr_context,
46 ast.Expression,
47 ast.For,
48 ast.Gt,
49 ast.GtE,
50 ast.In,
51 ast.Is,
52 ast.IsNot,
53 ast.List,
54 ast.ListComp,
55 ast.Load,
56 ast.LShift,
57 ast.Lt,
58 ast.LtE,
59 ast.Mod,
60 ast.Name,
61 ast.Not,
62 ast.NotEq,
63 ast.NotIn, # codespell:ignore notin
64 ast.operator,
65 ast.Or,
66 ast.RShift,
67 ast.Set,
68 ast.Slice,
69 ast.slice,
70 ast.Sub,
71 ast.Subscript,
72 ast.Tuple,
73 ast.UAdd,
74 ast.UnaryOp,
75 ast.unaryop,
76 ast.USub,
77)
80class DatabooksParser(ast.NodeVisitor):
81 """AST parser that disallows unsafe nodes/values."""
83 def __init__(self, **variables: Any) -> None:
84 """Instantiate with variables and callables (built-ins) scope."""
85 self.builtins = {b.__name__: b for b in _ALLOWED_BUILTINS}
86 self.names = deepcopy(variables) or {}
87 self.scope = {
88 **self.names,
89 "__builtins__": self.builtins,
90 }
92 @staticmethod
93 def _prioritize(field: Tuple[str, Any]) -> bool:
94 """Prioritize `ast.comprehension` nodes when expanding the AST tree."""
95 _, value = field
96 if not isinstance(value, list):
97 return True
98 return not any(isinstance(f, ast.comprehension) for f in value)
100 @staticmethod
101 def _allowed_attr(obj: Any, attr: str, is_dynamic: bool = False) -> None:
102 """
103 Check that attribute is a key of `databooks.data_models.base.DatabooksBase`.
105 If `obj` is an iterable and was computed dynamically (that is, not originally in
106 scope but computed from a comprehension), check attributes for all elements in
107 the iterable.
108 """
109 allowed_attrs = list(dict(obj).keys()) if isinstance(obj, DatabooksBase) else ()
110 if isinstance(obj, abc.Iterable) and is_dynamic:
111 for el in obj:
112 DatabooksParser._allowed_attr(obj=el, attr=attr)
113 else:
114 if attr not in allowed_attrs:
115 raise ValueError(
116 "Expected attribute to be one of"
117 f" `{allowed_attrs}`, got `{attr}` for {obj}."
118 )
120 def _get_iter(self, node: ast.expr) -> Iterable:
121 """Use `DatabooksParser.safe_eval_ast` to get the iterable object."""
122 tree = ast.Expression(body=node)
123 return iter(self.safe_eval_ast(tree))
125 def generic_visit(self, node: ast.AST) -> None:
126 """
127 Prioritize `ast.comprehension` nodes when expanding tree.
129 Similar to `NodeVisitor.generic_visit`, but favor comprehensions when multiple
130 nodes on the same level. In comprehensions, we have a generator argument that
131 includes names that are stored. By visiting them first we avoid 'running into'
132 unknown names.
133 """
134 if not isinstance(node, _ALLOWED_NODES):
135 raise ValueError(f"Invalid node `{node}`.")
137 for field, value in sorted(ast.iter_fields(node), key=self._prioritize):
138 if isinstance(value, list):
139 for item in value:
140 if isinstance(item, ast.AST):
141 self.visit(item)
142 elif isinstance(value, ast.AST):
143 self.visit(value)
145 def visit_comprehension(self, node: ast.comprehension) -> None:
146 """Add variable from a comprehension to list of allowed names."""
147 if not isinstance(node.target, ast.Name):
148 raise RuntimeError(
149 "Expected `ast.comprehension`'s target to be `ast.Name`, got"
150 f" `ast.{type(node.target).__name__}`."
151 )
152 self.names[node.target.id] = self._get_iter(node.iter)
153 self.generic_visit(node)
155 def visit_Attribute(self, node: ast.Attribute) -> None:
156 """Allow attributes for Pydantic fields only."""
157 if not isinstance(node.value, (ast.Attribute, ast.Name, ast.Subscript)):
158 raise ValueError(
159 "Expected attribute to be one of `ast.Name`, `ast.Attribute` or"
160 f" `ast.Subscript`, got `ast.{type(node.value).__name__}`."
161 )
162 if isinstance(node.value, ast.Name):
163 self._allowed_attr(
164 obj=self.names[node.value.id],
165 attr=node.attr,
166 is_dynamic=node.value.id in (self.names.keys() - self.scope.keys()),
167 )
168 self.generic_visit(node)
170 def visit_Name(self, node: ast.Name) -> None:
171 """Only allow names from scope or comprehension variables."""
172 valid_names = {**self.names, **self.builtins}
173 if node.id not in valid_names:
174 raise ValueError(
175 f"Expected `name` to be one of `{valid_names.keys()}`, got `{node.id}`."
176 )
177 self.generic_visit(node)
179 def safe_eval_ast(self, ast_tree: ast.Expression) -> Any:
180 """Evaluate safe AST trees only (raise errors otherwise)."""
181 self.visit(ast_tree)
182 exe = compile(ast_tree, filename="", mode="eval")
183 return eval(exe, self.scope)
185 def safe_eval(self, src: str) -> Any:
186 """
187 Evaluate strings that are safe only (raise errors otherwise).
189 A "safe" string or node provided may only consist of nodes in
190 `databooks.affirm._ALLOWED_NODES` and built-ins from
191 `databooks.affirm._ALLOWED_BUILTINS`.
192 """
193 ast_tree = ast.parse(src, mode="eval")
194 return self.safe_eval_ast(ast_tree)
197def affirm(nb_path: Path, exprs: List[str], verbose: bool = False) -> bool:
198 """
199 Return whether notebook passed all checks (expressions).
201 :param nb_path: Path of notebook file
202 :param exprs: Expression with check to be evaluated on notebook
203 :param verbose: Log failed tests for notebook
204 :return: Evaluated expression cast as a `bool`
205 """
206 if verbose:
207 set_verbose(logger)
209 nb = JupyterNotebook.parse_file(nb_path)
210 variables: Dict[str, Any] = {
211 "nb": nb,
212 "raw_cells": [c for c in nb.cells if c.cell_type == "raw"],
213 "md_cells": [c for c in nb.cells if c.cell_type == "markdown"],
214 "code_cells": [c for c in nb.cells if c.cell_type == "code"],
215 "exec_cells": [
216 c
217 for c in nb.cells
218 if c.cell_type == "code" and c.execution_count is not None
219 ],
220 }
221 databooks_parser = DatabooksParser(**variables)
222 is_ok = [bool(databooks_parser.safe_eval(expr)) for expr in exprs]
223 n_fail = sum([not ok for ok in is_ok])
225 logger.info(f"{nb_path} failed {n_fail} of {len(is_ok)} checks.")
226 logger.debug(
227 str(nb_path)
228 + (
229 f" failed {list(compress(exprs, (not ok for ok in is_ok)))}."
230 if n_fail > 0
231 else " succeeded all checks."
232 )
233 )
234 return all(is_ok)
237def affirm_all(
238 nb_paths: List[Path],
239 *,
240 progress_callback: Callable[[], None] = lambda: None,
241 **affirm_kwargs: Any,
242) -> List[bool]:
243 """
244 Clear metadata for multiple notebooks at notebooks and cell level.
246 :param nb_paths: Paths of notebooks to assert metadata
247 :param progress_callback: Callback function to report progress
248 :param affirm_kwargs: Keyword arguments to be passed to `databooks.affirm.affirm`
249 :return: Whether the notebooks contained or not the desired metadata
250 """
251 checks = []
252 for nb_path in nb_paths:
253 checks.append(affirm(nb_path, **affirm_kwargs))
254 progress_callback()
255 return checks