sqllocks-shape-behavior 0.9.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,40 @@
1
+ """Shape plugin: behavior models (state machines on a virtual clock).
2
+
3
+ The public names are importable from here, lazily: importing the package loads no module of
4
+ this plugin beyond this file, and none of NumPy or Arrow. See ``docs/plugins/behavior.md``.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import importlib
10
+ from typing import Any
11
+
12
+ SHAPE_API = "1.0"
13
+
14
+ _EXPORTS = {
15
+ "Module": "shape_behavior.model",
16
+ "ModuleError": "shape_behavior.model",
17
+ "load_module": "shape_behavior.model",
18
+ "Population": "shape_behavior.population",
19
+ "SimConfig": "shape_behavior.simulator",
20
+ "Simulator": "shape_behavior.simulator",
21
+ "Checkpoint": "shape_behavior.simulator",
22
+ "EVENT_SCHEMA": "shape_behavior.events",
23
+ "import_gmf": "shape_behavior.gmf",
24
+ "ImportResult": "shape_behavior.gmf",
25
+ "Unsupported": "shape_behavior.gmf",
26
+ "UnsupportedGmfError": "shape_behavior.gmf",
27
+ "register_state_type": "shape_behavior.extension",
28
+ "StateHandler": "shape_behavior.extension",
29
+ "StateContext": "shape_behavior.extension",
30
+ "Emission": "shape_behavior.extension",
31
+ }
32
+
33
+ __all__ = sorted(_EXPORTS)
34
+
35
+
36
+ def __getattr__(name: str) -> Any:
37
+ module = _EXPORTS.get(name)
38
+ if module is None:
39
+ raise AttributeError(f"module 'shape_behavior' has no attribute {name!r}")
40
+ return getattr(importlib.import_module(module), name)
@@ -0,0 +1,150 @@
1
+ """Behavior modules as ``shape.behaviors`` plugins (plugin API v1, ``Behavior`` protocol).
2
+
3
+ A :class:`ModuleBehavior` wraps one :class:`~shape_behavior.model.Module` so the engine, the
4
+ plugin host, ``shape plugins list`` and the conformance kit can all reach it by name::
5
+
6
+ from shape_behavior.behaviors import behavior
7
+
8
+ b = behavior({"format": "shape-behavior/1", "name": "my_process", "states": {...}})
9
+ events = b.simulate(population=1000, seed=7, years=3) # a pyarrow.Table
10
+
11
+ The built-in examples register under ``shape.behaviors`` through the classes at the bottom of
12
+ this file; a third-party package registers its own module the same way
13
+ (``docs/plugins/behavior.md``, section 10).
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from typing import Any
19
+
20
+ from shape_behavior.extension import get_handler
21
+ from shape_behavior.model import EVENT_KINDS, Module, load_module
22
+
23
+ SHAPE_API = "1.0"
24
+ VERSION = "1.0"
25
+
26
+ DEFAULT_POPULATION: dict[str, Any] = {
27
+ "age_at_start": {"kind": "uniform", "low": 0, "high": 90},
28
+ "attributes": {"gender": {"kind": "categorical", "values": {"F": 0.5, "M": 0.5}}},
29
+ }
30
+
31
+
32
+ def population_settings(modules: list[Module]) -> dict[str, Any]:
33
+ """The population settings the modules ask for (``population_defaults``; the first module
34
+ wins on a clash), or :data:`DEFAULT_POPULATION` when none asks."""
35
+ import copy
36
+
37
+ spec: dict[str, Any] = {}
38
+ for m in reversed(modules):
39
+ spec.update(m.doc.get("population_defaults", {}))
40
+ return spec or copy.deepcopy(DEFAULT_POPULATION)
41
+
42
+
43
+ class ModuleBehavior:
44
+ """One module as a ``shape.behaviors`` plugin: name, version, states, attributes, events."""
45
+
46
+ def __init__(self, module: Module | dict[str, Any] | str, version: str = VERSION) -> None:
47
+ self.module = load_module(module)
48
+ self.name: str = self.module.name
49
+ self.version: str = str(self.module.doc.get("version", version))
50
+ self.states: list[str] = list(self.module.states)
51
+ self.attributes: list[str] = sorted(self.module.attribute_uses())
52
+ self.events: list[str] = _event_kinds(self.module)
53
+
54
+ def simulate(self, population: int, seed: int, years: float) -> Any:
55
+ """Events of ``population`` entities over ``years`` from 2020-01-01, deterministic per
56
+ ``seed`` (an Arrow table in :data:`shape_behavior.EVENT_SCHEMA`)."""
57
+ from shape_behavior.population import Population
58
+ from shape_behavior.simulator import SimConfig, Simulator
59
+ from shape_behavior.timeutil import add_years, to_us
60
+
61
+ start = "2020-01-01"
62
+ pop = Population.from_dict(population_settings([self.module]), size=population, start=start)
63
+ sim = Simulator([self.module], pop, SimConfig(seed=seed))
64
+ return sim.run_until(add_years(to_us(start), years))
65
+
66
+
67
+ def behavior(module: Module | dict[str, Any] | str, version: str = VERSION) -> ModuleBehavior:
68
+ """A :class:`ModuleBehavior` for a module, a module document or a module file path."""
69
+ return ModuleBehavior(module, version)
70
+
71
+
72
+ def _event_kinds(module: Module) -> list[str]:
73
+ kinds: set[str] = set()
74
+ for state in module.states.values():
75
+ kind = state["type"]
76
+ if kind == "event":
77
+ kinds.add(str(state["event"]))
78
+ elif kind in EVENT_KINDS:
79
+ kinds.add(EVENT_KINDS[kind])
80
+ else: # a registered state type may declare the kinds it emits (optional `kinds`)
81
+ declared = getattr(get_handler(kind), "kinds", None)
82
+ if declared is not None:
83
+ kinds.update(declared(state))
84
+ return sorted(kinds)
85
+
86
+
87
+ class Subscription(ModuleBehavior):
88
+ """Subscription lifecycle: free trial, conversion, monthly payments, pause and cancel."""
89
+
90
+ def __init__(self) -> None:
91
+ super().__init__("subscription")
92
+
93
+
94
+ class EquipmentMaintenance(ModuleBehavior):
95
+ """Equipment maintenance: run, degrade, fail, repair, with scheduled inspections."""
96
+
97
+ def __init__(self) -> None:
98
+ super().__init__("equipment_maintenance")
99
+
100
+
101
+ class HealthcareScreening(ModuleBehavior):
102
+ """A tiny healthcare example (screening, diagnosis, treatment), written from scratch."""
103
+
104
+ def __init__(self) -> None:
105
+ super().__init__("healthcare_screening")
106
+
107
+
108
+ class EventSequence(ModuleBehavior):
109
+ """An ordered funnel of named events with a continue probability per step (defaults)."""
110
+
111
+ def __init__(self) -> None:
112
+ from shape_behavior.primitives import event_sequence
113
+
114
+ super().__init__(event_sequence())
115
+
116
+
117
+ class TelemetrySeries(ModuleBehavior):
118
+ """A regular reading per device with drift, noise, skipped and stuck readings (defaults)."""
119
+
120
+ def __init__(self) -> None:
121
+ from shape_behavior.primitives import telemetry_series
122
+
123
+ super().__init__(telemetry_series())
124
+
125
+
126
+ class TransactionStream(ModuleBehavior):
127
+ """Transactions as a Poisson process per entity, with refunds and reversals (defaults)."""
128
+
129
+ def __init__(self) -> None:
130
+ from shape_behavior.primitives import transaction_stream
131
+
132
+ super().__init__(transaction_stream())
133
+
134
+
135
+ class FileArrival(ModuleBehavior):
136
+ """A feed with an expected arrival per schedule slot: on time, late, missing, duplicated."""
137
+
138
+ def __init__(self) -> None:
139
+ from shape_behavior.primitives import file_arrival
140
+
141
+ super().__init__(file_arrival())
142
+
143
+
144
+ class EntityLifecycle(ModuleBehavior):
145
+ """Create, update and delete events per entity (change-data-capture style)."""
146
+
147
+ def __init__(self) -> None:
148
+ from shape_behavior.primitives import entity_lifecycle
149
+
150
+ super().__init__(entity_lifecycle())
shape_behavior/cli.py ADDED
@@ -0,0 +1,240 @@
1
+ """``shape behave``: run, check and import behavior modules.
2
+
3
+ Reference: ``docs/plugins/behavior.md``, section 7.
4
+
5
+ ::
6
+
7
+ shape behave run subscription --population 10000 --years 3 --seed 7 -o out/
8
+ shape behave check my_module.json
9
+ shape behave import-gmf downloaded_module.json -o converted.json
10
+ shape behave examples -o modules/
11
+
12
+ Nothing heavy loads at import time: NumPy, Arrow and the simulator load when a command runs.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ import sys
19
+ import time
20
+ from pathlib import Path
21
+ from typing import Any
22
+
23
+ from shape_behavior.params import (
24
+ PRIMITIVE_NAMES,
25
+ PrimitiveParamsError,
26
+ read_params,
27
+ unknown_primitive,
28
+ )
29
+
30
+ SHAPE_API = "1.0"
31
+
32
+
33
+ class BehaveCommand:
34
+ """``shape behave COMMAND ...``."""
35
+
36
+ name = "behave"
37
+ help = "run state-machine behavior modules on a virtual clock (events as Parquet)"
38
+
39
+ def configure(self, parser: Any) -> None:
40
+ sub = parser.add_subparsers(dest="behave_cmd", required=True, metavar="COMMAND")
41
+ run = sub.add_parser("run", help="simulate a population and write the event stream")
42
+ run.add_argument(
43
+ "modules",
44
+ nargs="+",
45
+ metavar="MODULES",
46
+ help="module files (native or GMF JSON) or example names",
47
+ )
48
+ run.add_argument(
49
+ "--population", type=int, required=True, metavar="N", help="number of entities"
50
+ )
51
+ run.add_argument(
52
+ "--years", type=float, required=True, metavar="Y", help="years to simulate"
53
+ )
54
+ run.add_argument("--seed", type=int, default=0, help="random seed (default 0)")
55
+ run.add_argument("-o", "--out", required=True, metavar="OUT", help="output directory")
56
+ run.add_argument(
57
+ "--start", default="2020-01-01", help="virtual clock start (default 2020-01-01)"
58
+ )
59
+ run.add_argument(
60
+ "--population-spec",
61
+ metavar="FILE",
62
+ help="JSON population settings (attributes, age_at_start, arrival, lifetime)",
63
+ )
64
+ run.add_argument(
65
+ "--window-years", type=float, default=1.0, help="years per output file (default 1)"
66
+ )
67
+ run.add_argument(
68
+ "--poll", default="7 days", help="guard polling interval (default '7 days')"
69
+ )
70
+ run.add_argument("--strict", action="store_true", help="fail on unsupported GMF elements")
71
+ run.add_argument(
72
+ "--params",
73
+ metavar="FILE.json",
74
+ help="parameters of a primitive named in MODULES "
75
+ '({"format": "shape-behavior-params", "version": 1, "primitive": NAME, "params": {}})',
76
+ )
77
+ run.set_defaults(fn=_run)
78
+ check = sub.add_parser("check", help="validate modules; exit 1 when one is invalid")
79
+ check.add_argument("modules", nargs="+", metavar="MODULES")
80
+ check.add_argument("--strict", action="store_true", help="fail on unsupported GMF elements")
81
+ check.set_defaults(fn=_check)
82
+ imp = sub.add_parser(
83
+ "import-gmf", help="convert a GMF module and report what is unsupported"
84
+ )
85
+ imp.add_argument("file", help="the GMF JSON file you downloaded")
86
+ imp.add_argument("-o", "--out", metavar="OUT.json", help="write the native module here")
87
+ imp.add_argument("--strict", action="store_true", help="fail on unsupported elements")
88
+ imp.set_defaults(fn=_import)
89
+ ex = sub.add_parser("examples", help="list the built-in example modules, or write them")
90
+ ex.add_argument("-o", "--out", metavar="DIR", help="write the examples into DIR")
91
+ ex.set_defaults(fn=_examples)
92
+
93
+ def run(self, args: Any) -> int:
94
+ try:
95
+ return int(args.fn(args))
96
+ except (FileNotFoundError, json.JSONDecodeError, PrimitiveParamsError) as exc:
97
+ print(f"shape behave: {exc}", file=sys.stderr)
98
+ return 2
99
+ except ValueError as exc: # includes ModuleError, UnsupportedGmfError
100
+ print(f"shape behave: {exc}", file=sys.stderr)
101
+ return 1
102
+
103
+
104
+ def _load(sources: list[str], strict: bool, params: str | None = None) -> list[Any]:
105
+ from shape_behavior.model import load_module
106
+
107
+ primitive, values = read_params(Path(params)) if params else ("", {})
108
+ if params and primitive not in sources:
109
+ if primitive not in PRIMITIVE_NAMES:
110
+ raise unknown_primitive(primitive)
111
+ raise PrimitiveParamsError(
112
+ f"--params is for primitive {primitive!r}, which is not among the modules to run "
113
+ f"({', '.join(sources)})",
114
+ "primitive",
115
+ )
116
+ if params:
117
+ from shape_behavior.primitives import build
118
+ modules = []
119
+ for src in sources:
120
+ module = (
121
+ build(primitive, values)
122
+ if params and src == primitive
123
+ else load_module(src, strict=strict)
124
+ )
125
+ report = module.import_report
126
+ if report is not None and (report.unsupported or report.warnings):
127
+ print(f"{src}:\n{report.report()}", file=sys.stderr)
128
+ modules.append(module)
129
+ return modules
130
+
131
+
132
+ def _run(args: Any) -> int:
133
+ import pyarrow.parquet as pq # type: ignore[import-untyped]
134
+
135
+ from shape_behavior.behaviors import population_settings
136
+ from shape_behavior.population import Population
137
+ from shape_behavior.simulator import SimConfig, Simulator
138
+ from shape_behavior.timeutil import add_years, to_us
139
+
140
+ if args.population < 0 or args.years <= 0 or args.window_years <= 0:
141
+ print(
142
+ "shape behave: --population must be >= 0 and --years and --window-years > 0",
143
+ file=sys.stderr,
144
+ )
145
+ return 2
146
+ modules = _load(args.modules, args.strict, args.params)
147
+ if args.population_spec:
148
+ spec = json.loads(Path(args.population_spec).read_text(encoding="utf-8"))
149
+ else:
150
+ spec = population_settings(modules)
151
+ population = Population.from_dict(spec, size=args.population, start=args.start)
152
+ sim = Simulator(modules, population, SimConfig(seed=args.seed, poll=args.poll))
153
+ out = Path(args.out)
154
+ (out / "events").mkdir(parents=True, exist_ok=True)
155
+ start = to_us(args.start)
156
+ end = add_years(start, args.years)
157
+ started = time.perf_counter()
158
+ windows: list[dict[str, Any]] = []
159
+ cursor, part = start, 0
160
+ while cursor < end:
161
+ nxt = min(add_years(start, args.window_years * (part + 1)), end)
162
+ events = sim.run_until(nxt)
163
+ pq.write_table(events, out / "events" / f"part-{part:04d}.parquet", compression="snappy")
164
+ windows.append(
165
+ {
166
+ "file": f"events/part-{part:04d}.parquet",
167
+ "until": str(sim.now),
168
+ "rows": events.num_rows,
169
+ }
170
+ )
171
+ cursor, part = nxt, part + 1
172
+ pq.write_table(sim.entities(), out / "entities.parquet", compression="snappy")
173
+ elapsed = time.perf_counter() - started
174
+ total = sum(w["rows"] for w in windows)
175
+ manifest = {
176
+ "seed": args.seed,
177
+ "years": args.years,
178
+ "start": args.start,
179
+ "population": population.to_dict(),
180
+ "modules": [_module_entry(m, s) for m, s in zip(modules, args.modules, strict=True)],
181
+ "events": total,
182
+ "windows": windows,
183
+ "seconds": round(elapsed, 3),
184
+ "unsupported": [
185
+ str(u) for m in modules if m.import_report for u in m.import_report.unsupported
186
+ ],
187
+ }
188
+ (out / "run.json").write_text(
189
+ json.dumps(manifest, indent=2) + "\n", encoding="utf-8", newline="\n"
190
+ )
191
+ print(f"{total:,} events for {args.population:,} entities in {elapsed:.1f}s -> {out}")
192
+ return 0
193
+
194
+
195
+ def _module_entry(module: Any, source: str) -> dict[str, Any]:
196
+ entry: dict[str, Any] = {"name": module.name, "digest": module.digest(), "source": source}
197
+ if "primitive" in module.doc: # a primitive: record the parameters it ran with
198
+ entry["primitive"] = module.doc["primitive"]
199
+ entry["parameters"] = module.doc["parameters"]
200
+ return entry
201
+
202
+
203
+ def _check(args: Any) -> int:
204
+ modules = _load(args.modules, args.strict)
205
+ for m in modules:
206
+ print(f"ok {m.name} ({len(m.states)} states)")
207
+ return 0
208
+
209
+
210
+ def _import(args: Any) -> int:
211
+ from shape_behavior.gmf import import_gmf
212
+
213
+ result = import_gmf(Path(args.file), strict=args.strict)
214
+ print(result.report())
215
+ if args.out:
216
+ Path(args.out).write_text(
217
+ json.dumps(result.module.to_dict(), indent=2) + "\n", encoding="utf-8", newline="\n"
218
+ )
219
+ print(f"wrote {args.out}")
220
+ return 0
221
+
222
+
223
+ def _examples(args: Any) -> int:
224
+ import importlib.resources
225
+
226
+ from shape_behavior.model import examples
227
+
228
+ names = examples()
229
+ if not args.out:
230
+ print("\n".join(names))
231
+ return 0
232
+ out = Path(args.out)
233
+ out.mkdir(parents=True, exist_ok=True)
234
+ for name in names:
235
+ text = (
236
+ importlib.resources.files("shape_behavior") / "examples" / f"{name}.json"
237
+ ).read_text(encoding="utf-8")
238
+ (out / f"{name}.json").write_text(text, encoding="utf-8", newline="\n")
239
+ print(f"wrote {len(names)} modules to {out}")
240
+ return 0
@@ -0,0 +1,221 @@
1
+ """Conditions: validation, and compilation into vectorized tests over rows of entities."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Callable
6
+ from typing import Any
7
+
8
+ import numpy as np
9
+
10
+ from shape_behavior.timeutil import to_us, unit_us
11
+
12
+ OPS = ("==", "!=", "<", "<=", ">", ">=")
13
+ NIL_OPS = ("is_nil", "is_not_nil")
14
+ LEAF_TYPES = ("true", "false", "attribute", "age", "date", "active_condition", "active_medication")
15
+ GROUP_TYPES = ("and", "or", "at_least", "at_most")
16
+ _CMP: dict[str, Callable[[Any, Any], Any]] = {
17
+ "==": np.equal,
18
+ "!=": np.not_equal,
19
+ "<": np.less,
20
+ "<=": np.less_equal,
21
+ ">": np.greater,
22
+ ">=": np.greater_equal,
23
+ }
24
+ _Test = Callable[[Any, Any, Any], Any]
25
+
26
+
27
+ def _num(v: Any) -> float | None:
28
+ if isinstance(v, bool):
29
+ return 1.0 if v else 0.0
30
+ if isinstance(v, (int, float)):
31
+ return float(v)
32
+ return None
33
+
34
+
35
+ def check(doc: Any, where: str) -> list[str]:
36
+ """Problems with a condition document, empty when it is fine."""
37
+ if not isinstance(doc, dict):
38
+ return [f"{where}: a condition must be an object"]
39
+ kind = doc.get("type")
40
+ if kind in ("true", "false"):
41
+ return []
42
+ if kind == "attribute":
43
+ out = [] if isinstance(doc.get("attribute"), str) else [f"{where}: needs 'attribute'"]
44
+ op = doc.get("op")
45
+ if op in NIL_OPS:
46
+ return out
47
+ if op not in OPS:
48
+ return [*out, f"{where}: op must be one of {', '.join(OPS + NIL_OPS)}"]
49
+ if "value" not in doc:
50
+ out.append(f"{where}: needs 'value'")
51
+ elif isinstance(doc["value"], str) and op not in ("==", "!="):
52
+ out.append(f"{where}: text attributes only compare with == and !=")
53
+ elif _num(doc["value"]) is None and not isinstance(doc["value"], str):
54
+ out.append(f"{where}: value must be a number, a bool or text")
55
+ return out
56
+ if kind in ("age", "date"):
57
+ out = []
58
+ if doc.get("op") not in OPS:
59
+ out.append(f"{where}: op must be one of {', '.join(OPS)}")
60
+ if "value" not in doc:
61
+ return [*out, f"{where}: needs 'value'"]
62
+ try:
63
+ if kind == "age":
64
+ if _num(doc["value"]) is None:
65
+ raise ValueError("age value must be a number")
66
+ unit_us(doc.get("unit", "years"))
67
+ else:
68
+ to_us(doc["value"])
69
+ except (ValueError, TypeError) as exc:
70
+ out.append(f"{where}: {exc}")
71
+ return out
72
+ if kind in ("active_condition", "active_medication"):
73
+ codes = doc.get("codes")
74
+ if not isinstance(codes, list) or not codes or not all(code_of(c) for c in codes):
75
+ return [f"{where}: needs a non-empty list of 'codes'"]
76
+ return []
77
+ if kind in GROUP_TYPES:
78
+ subs = doc.get("conditions")
79
+ if not isinstance(subs, list) or not subs:
80
+ return [f"{where}: needs a non-empty list of 'conditions'"]
81
+ out = []
82
+ if kind in ("at_least", "at_most"):
83
+ key = "minimum" if kind == "at_least" else "maximum"
84
+ if isinstance(doc.get(key), bool) or not isinstance(doc.get(key), int):
85
+ out.append(f"{where}: needs an integer {key!r}")
86
+ for i, sub in enumerate(subs):
87
+ out += check(sub, f"{where}.conditions[{i}]")
88
+ return out
89
+ if kind == "not":
90
+ return check(doc.get("condition"), f"{where}.condition")
91
+ return [f"{where}: unknown condition type {kind!r}"]
92
+
93
+
94
+ def code_of(c: Any) -> str | None:
95
+ """The code string of a ``codes`` entry (a dict with ``code``, or the string itself)."""
96
+ if isinstance(c, str) and c:
97
+ return c
98
+ if isinstance(c, dict) and isinstance(c.get("code"), str) and c["code"]:
99
+ return str(c["code"])
100
+ return None
101
+
102
+
103
+ def walk(doc: Any) -> list[dict[str, Any]]:
104
+ """Every condition node in ``doc``, the root first."""
105
+ out = [doc]
106
+ if doc.get("type") in GROUP_TYPES:
107
+ for sub in doc["conditions"]:
108
+ out += walk(sub)
109
+ elif doc.get("type") == "not":
110
+ out += walk(doc["condition"])
111
+ return out
112
+
113
+
114
+ class Compiled:
115
+ """A compiled condition: ``test(store, rows, time)`` gives a bool array over the rows;
116
+ ``ready`` (``None`` unless the condition depends only on age) gives the earliest time at
117
+ which it holds, ``INF`` for never."""
118
+
119
+ def __init__(self, test: _Test, ready: Callable[[Any, Any], Any] | None = None) -> None:
120
+ self.test = test
121
+ self.ready = ready
122
+
123
+
124
+ def compile_condition(doc: dict[str, Any], bit: Callable[[str, str], int]) -> Compiled:
125
+ """Compile ``doc``; ``bit(kind, code)`` gives the bit of a code in the ``condition`` or
126
+ ``medication`` active set."""
127
+ kind = doc["type"]
128
+ if kind == "true":
129
+ return Compiled(lambda s, r, t: np.ones(len(r), bool))
130
+ if kind == "false":
131
+ return Compiled(lambda s, r, t: np.zeros(len(r), bool))
132
+ if kind == "attribute":
133
+ return Compiled(_attribute_test(doc))
134
+ if kind == "age":
135
+ return _age(doc)
136
+ if kind == "date":
137
+ value = to_us(doc["value"])
138
+ cmp = _CMP[doc["op"]]
139
+ return Compiled(lambda s, r, t: cmp(t, value))
140
+ if kind in ("active_condition", "active_medication"):
141
+ mask = 0
142
+ which = "condition" if kind == "active_condition" else "medication"
143
+ for c in doc["codes"]:
144
+ mask |= 1 << bit(which, str(code_of(c)))
145
+ field = "cond_mask" if which == "condition" else "med_mask"
146
+ m = np.uint64(mask)
147
+ return Compiled(lambda s, r, t: (getattr(s, field)[r] & m) != 0)
148
+ if kind == "not":
149
+ inner = compile_condition(doc["condition"], bit).test
150
+ return Compiled(lambda s, r, t: ~inner(s, r, t))
151
+ subs = [compile_condition(c, bit) for c in doc["conditions"]]
152
+ tests = [c.test for c in subs]
153
+ if kind == "and":
154
+ ready = _and_ready(subs)
155
+ return Compiled(lambda s, r, t: np.logical_and.reduce([f(s, r, t) for f in tests]), ready)
156
+ if kind == "or":
157
+ return Compiled(lambda s, r, t: np.logical_or.reduce([f(s, r, t) for f in tests]))
158
+ count = doc.get("minimum" if kind == "at_least" else "maximum", 0)
159
+
160
+ def counted(s: Any, r: Any, t: Any) -> Any:
161
+ n = np.sum([f(s, r, t) for f in tests], axis=0)
162
+ return n >= count if kind == "at_least" else n <= count
163
+
164
+ return Compiled(counted)
165
+
166
+
167
+ def _and_ready(subs: list[Compiled]) -> Callable[[Any, Any], Any] | None:
168
+ readies = [c.ready for c in subs]
169
+ if any(r is None for r in readies):
170
+ return None
171
+ fns: list[Callable[[Any, Any], Any]] = [r for r in readies if r is not None]
172
+ return lambda s, r: np.maximum.reduce([f(s, r) for f in fns])
173
+
174
+
175
+ def _attribute_test(doc: dict[str, Any]) -> _Test:
176
+ name, op = doc["attribute"], doc["op"]
177
+ if op == "is_nil":
178
+ return lambda s, r, t: s.is_nil(name, r)
179
+ if op == "is_not_nil":
180
+ return lambda s, r, t: ~s.is_nil(name, r)
181
+ value = doc["value"]
182
+ cmp = _CMP[op]
183
+ if isinstance(value, str):
184
+
185
+ def text_test(s: Any, r: Any, t: Any) -> Any:
186
+ if name in s.num: # a numeric attribute never equals text
187
+ return np.full(len(r), op == "!=")
188
+ code = s.vindex.get(name, {}).get(value, -2)
189
+ return cmp(s.cat[name][r], code) if name in s.cat else np.full(len(r), op == "!=")
190
+
191
+ return text_test
192
+ num = float(_num(value) or 0.0)
193
+
194
+ def num_test(s: Any, r: Any, t: Any) -> Any:
195
+ if name not in s.num:
196
+ return np.zeros(len(r), bool)
197
+ x = s.num[name][r]
198
+ res = cmp(x, num)
199
+ return res | np.isnan(x) if op == "!=" else res
200
+
201
+ return num_test
202
+
203
+
204
+ def _age(doc: dict[str, Any]) -> Compiled:
205
+ threshold = round(float(doc["value"]) * unit_us(doc.get("unit", "years")))
206
+ op = doc["op"]
207
+ cmp = _CMP[op]
208
+
209
+ def test(s: Any, r: Any, t: Any) -> Any:
210
+ return cmp(t - s.born[r], threshold)
211
+
212
+ ready: Callable[[Any, Any], Any] | None = None
213
+ if op in (">=", ">"):
214
+ # the first microsecond at which age >= value (or > value)
215
+ def ready(s: Any, r: Any) -> Any:
216
+ return s.born[r] + threshold + (1 if op == ">" else 0)
217
+
218
+ elif op in ("<", "<="):
219
+ # holds from the start and stops: handled by the caller as "now or never"
220
+ ready = None
221
+ return Compiled(test, ready)