lf-pollywog 0.1.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.
- docs/conf.py +39 -0
- lf_pollywog-0.1.1.dist-info/METADATA +166 -0
- lf_pollywog-0.1.1.dist-info/RECORD +22 -0
- lf_pollywog-0.1.1.dist-info/WHEEL +5 -0
- lf_pollywog-0.1.1.dist-info/licenses/LICENSE +21 -0
- lf_pollywog-0.1.1.dist-info/top_level.txt +3 -0
- pollywog/__init__.py +2 -0
- pollywog/conversion/__init__.py +0 -0
- pollywog/conversion/sklearn.py +154 -0
- pollywog/core.py +838 -0
- pollywog/display.py +149 -0
- pollywog/helpers.py +246 -0
- pollywog/run.py +116 -0
- pollywog/utils.py +95 -0
- tests/__init__.py +0 -0
- tests/conftest.py +4 -0
- tests/test_conversion.py +72 -0
- tests/test_core.py +274 -0
- tests/test_display.py +60 -0
- tests/test_helpers.py +72 -0
- tests/test_run.py +71 -0
- tests/test_utils.py +32 -0
pollywog/core.py
ADDED
@@ -0,0 +1,838 @@
|
|
1
|
+
import json
|
2
|
+
import re
|
3
|
+
import zlib
|
4
|
+
from pathlib import Path
|
5
|
+
from typing import (
|
6
|
+
Any, Callable, Dict, List, Optional, Set, Type, TypeVar, Union, Sequence
|
7
|
+
)
|
8
|
+
|
9
|
+
from .utils import ensure_list, ensure_str_list, to_dict
|
10
|
+
|
11
|
+
HEADER = b"\x25\x6c\x66\x63\x61\x6c\x63\x2d\x31\x2e\x30\x0a\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
|
12
|
+
|
13
|
+
ITEM_ORDER = {
|
14
|
+
"variable": 0,
|
15
|
+
"calculation": 1,
|
16
|
+
"filter": 2,
|
17
|
+
}
|
18
|
+
|
19
|
+
|
20
|
+
# TODO: check if items need to be sorted into variables then calculations then filters and do so if needed
|
21
|
+
# TODO: actually just sorted preemptively when writing to file, will check later if this is an issue
|
22
|
+
class CalcSet:
|
23
|
+
"""
|
24
|
+
Represents a Leapfrog-style calculation set, containing variables, calculations, categories, filters, and conditional logic.
|
25
|
+
|
26
|
+
CalcSet is the main container for building, manipulating, and exporting calculation workflows. It is designed to help automate large, complex, or repetitive calculation sets, and supports querying, dependency analysis, and rich display in Jupyter notebooks.
|
27
|
+
|
28
|
+
Args:
|
29
|
+
items (list of Item): List of calculation items (Number, Category, Filter, If, etc.)
|
30
|
+
|
31
|
+
Example:
|
32
|
+
>>> from pollywog.core import CalcSet, Number
|
33
|
+
>>> calcset = CalcSet([
|
34
|
+
... Number(name="Au_est", children=["block[Au] * 0.95"]),
|
35
|
+
... Number(name="Ag_est", children=["block[Ag] * 0.85"])
|
36
|
+
... ])
|
37
|
+
"""
|
38
|
+
def __init__(self, items: Sequence[Union["Item", "Variable"]]):
|
39
|
+
"""
|
40
|
+
Initialize a CalcSet with a list of items.
|
41
|
+
Args:
|
42
|
+
items (list): List of calculation items (Number, Category, Filter, If, etc.)
|
43
|
+
"""
|
44
|
+
self.items = ensure_list(items)
|
45
|
+
|
46
|
+
def copy(self) -> "CalcSet":
|
47
|
+
"""
|
48
|
+
Return a deep copy of the CalcSet and its items.
|
49
|
+
"""
|
50
|
+
return CalcSet([item.copy() for item in self.items])
|
51
|
+
|
52
|
+
def query(self, expr: str, **external_vars) -> "CalcSet":
|
53
|
+
"""
|
54
|
+
Filter items in the CalcSet using a query expression.
|
55
|
+
|
56
|
+
The expression can use attributes of Item (e.g., 'item_type == "variable" and name.startswith("Au")'),
|
57
|
+
and external variables using @var syntax (like pandas).
|
58
|
+
|
59
|
+
Args:
|
60
|
+
expr (str): Query expression to filter items.
|
61
|
+
**external_vars: External variables to use in the query (referenced as @var).
|
62
|
+
|
63
|
+
Returns:
|
64
|
+
CalcSet: New CalcSet with filtered items.
|
65
|
+
|
66
|
+
Example:
|
67
|
+
>>> calcset.query('name.startswith("Au")')
|
68
|
+
"""
|
69
|
+
import inspect
|
70
|
+
import re
|
71
|
+
|
72
|
+
filtered = []
|
73
|
+
# Get caller's frame to access local and global variables
|
74
|
+
frame = inspect.currentframe()
|
75
|
+
try:
|
76
|
+
caller_frame = frame.f_back if frame is not None else None
|
77
|
+
caller_locals = caller_frame.f_locals if caller_frame else {}
|
78
|
+
caller_globals = caller_frame.f_globals if caller_frame else {}
|
79
|
+
finally:
|
80
|
+
del frame
|
81
|
+
|
82
|
+
# Merge external_vars with caller's scope, external_vars takes precedence
|
83
|
+
merged_vars = dict(caller_globals)
|
84
|
+
merged_vars.update(caller_locals)
|
85
|
+
merged_vars.update(external_vars)
|
86
|
+
|
87
|
+
# Safe helpers for query expressions
|
88
|
+
SAFE_EVAL_HELPERS = {
|
89
|
+
"len": len,
|
90
|
+
"any": any,
|
91
|
+
"all": all,
|
92
|
+
"min": min,
|
93
|
+
"max": max,
|
94
|
+
"sorted": sorted,
|
95
|
+
"re": re,
|
96
|
+
"str": str,
|
97
|
+
}
|
98
|
+
|
99
|
+
def replace_at_var(match):
|
100
|
+
var_name = match.group(1)
|
101
|
+
if var_name in merged_vars:
|
102
|
+
return f'merged_vars["{var_name}"]'
|
103
|
+
else:
|
104
|
+
raise NameError(f"External variable '@{var_name}' not provided.")
|
105
|
+
|
106
|
+
expr_eval = re.sub(r"@([A-Za-z_][A-Za-z0-9_]*)", replace_at_var, expr)
|
107
|
+
for item in self.items:
|
108
|
+
ns = {k: getattr(item, k, None) for k in dir(item) if not k.startswith("_")}
|
109
|
+
try:
|
110
|
+
if eval(
|
111
|
+
expr_eval, {"merged_vars": merged_vars, **SAFE_EVAL_HELPERS}, ns
|
112
|
+
):
|
113
|
+
filtered.append(item)
|
114
|
+
except Exception:
|
115
|
+
pass
|
116
|
+
return CalcSet(filtered)
|
117
|
+
|
118
|
+
def topological_sort(self) -> "CalcSet":
|
119
|
+
"""
|
120
|
+
Return a new CalcSet with items sorted topologically by dependencies.
|
121
|
+
|
122
|
+
This is useful for ensuring calculations are ordered so that dependencies are resolved before use.
|
123
|
+
Raises ValueError if cyclic dependencies are detected.
|
124
|
+
"""
|
125
|
+
items_by_name = {
|
126
|
+
item.name: item for item in self.items if hasattr(item, "name")
|
127
|
+
}
|
128
|
+
sorted_items = []
|
129
|
+
visited = set()
|
130
|
+
temp_mark = set()
|
131
|
+
|
132
|
+
def visit(item):
|
133
|
+
if item.name in visited:
|
134
|
+
return
|
135
|
+
if item.name in temp_mark:
|
136
|
+
raise ValueError(f"Cyclic dependency detected involving '{item.name}'")
|
137
|
+
temp_mark.add(item.name)
|
138
|
+
for dep in getattr(item, "dependencies", set()):
|
139
|
+
if dep in items_by_name:
|
140
|
+
visit(items_by_name[dep])
|
141
|
+
temp_mark.remove(item.name)
|
142
|
+
visited.add(item.name)
|
143
|
+
sorted_items.append(item)
|
144
|
+
|
145
|
+
for item in self.items:
|
146
|
+
visit(item)
|
147
|
+
|
148
|
+
# Add items without a name (should be rare)
|
149
|
+
unnamed = [item for item in self.items if not hasattr(item, "name")]
|
150
|
+
sorted_items.extend(unnamed)
|
151
|
+
|
152
|
+
return CalcSet(sorted_items)
|
153
|
+
|
154
|
+
# ...existing code...
|
155
|
+
"""
|
156
|
+
Return a deep copy of the CalcSet and its items.
|
157
|
+
"""
|
158
|
+
return CalcSet([item.copy() for item in self.items])
|
159
|
+
|
160
|
+
def to_dict(self, sort_items: bool = True) -> Dict[str, Any]:
|
161
|
+
"""
|
162
|
+
Convert the CalcSet to a dictionary representation.
|
163
|
+
|
164
|
+
Args:
|
165
|
+
sort_items (bool): Whether to sort items by type.
|
166
|
+
|
167
|
+
Returns:
|
168
|
+
dict: Dictionary representation of the calculation set.
|
169
|
+
"""
|
170
|
+
items = to_dict(self.items)
|
171
|
+
if sort_items:
|
172
|
+
items.sort(key=lambda x: ITEM_ORDER.get(x.get("type"), 99))
|
173
|
+
return {"type": "calculation-set", "items": items}
|
174
|
+
|
175
|
+
@classmethod
|
176
|
+
def from_dict(cls: Type["CalcSet"], data: Dict[str, Any]) -> "CalcSet":
|
177
|
+
"""
|
178
|
+
Create a CalcSet from a dictionary.
|
179
|
+
|
180
|
+
Args:
|
181
|
+
data (dict): Dictionary containing calculation set data.
|
182
|
+
|
183
|
+
Returns:
|
184
|
+
CalcSet: Instance of CalcSet.
|
185
|
+
"""
|
186
|
+
if data["type"] != "calculation-set":
|
187
|
+
raise ValueError(f"Expected type 'calculation-set', got {data['type']}")
|
188
|
+
items = []
|
189
|
+
for item in data["items"]:
|
190
|
+
item_type = item["type"]
|
191
|
+
if item_type in classes:
|
192
|
+
items.append(classes[item_type].from_dict(item))
|
193
|
+
elif item_type == "calculation":
|
194
|
+
if item.get("calculation_type") == "number":
|
195
|
+
items.append(Number.from_dict(item))
|
196
|
+
elif item.get("calculation_type") == "string":
|
197
|
+
items.append(Category.from_dict(item))
|
198
|
+
else:
|
199
|
+
raise ValueError(
|
200
|
+
f"Unknown calculation type: {item.get('calculation_type')}"
|
201
|
+
)
|
202
|
+
else:
|
203
|
+
raise ValueError(f"Unknown item type: {item_type}")
|
204
|
+
return cls(items=items)
|
205
|
+
|
206
|
+
def to_json(self, sort_items: bool = True, indent: int = 0) -> str:
|
207
|
+
"""
|
208
|
+
Convert the CalcSet to a JSON string.
|
209
|
+
|
210
|
+
Args:
|
211
|
+
sort_items (bool): Whether to sort items by type.
|
212
|
+
indent (int): Indentation level for JSON output.
|
213
|
+
|
214
|
+
Returns:
|
215
|
+
str: JSON string representation.
|
216
|
+
"""
|
217
|
+
return json.dumps(self.to_dict(sort_items=sort_items), indent=indent)
|
218
|
+
|
219
|
+
def to_lfcalc(
|
220
|
+
self, filepath_or_buffer: Union[str, Path, Any], sort_items: bool = True
|
221
|
+
) -> None:
|
222
|
+
"""
|
223
|
+
Write the CalcSet to a Leapfrog .lfcalc file.
|
224
|
+
|
225
|
+
Args:
|
226
|
+
filepath_or_buffer (str, Path, or file-like): Output file path or buffer.
|
227
|
+
sort_items (bool): Whether to sort items by type.
|
228
|
+
"""
|
229
|
+
if isinstance(filepath_or_buffer, (str, Path)):
|
230
|
+
with open(filepath_or_buffer, "wb") as f:
|
231
|
+
self._write_to_file(f, sort_items=sort_items)
|
232
|
+
else:
|
233
|
+
self._write_to_file(filepath_or_buffer, sort_items=sort_items)
|
234
|
+
|
235
|
+
def _write_to_file(self, file: Any, sort_items: bool) -> None:
|
236
|
+
"""
|
237
|
+
Write the CalcSet to a file in Leapfrog format.
|
238
|
+
Args:
|
239
|
+
file (file-like): File object to write to.
|
240
|
+
sort_items (bool): Whether to sort items by type.
|
241
|
+
"""
|
242
|
+
compressed_data = zlib.compress(
|
243
|
+
self.to_json(sort_items=sort_items).encode("utf-8")
|
244
|
+
)
|
245
|
+
file.write(HEADER)
|
246
|
+
file.write(compressed_data)
|
247
|
+
|
248
|
+
@staticmethod
|
249
|
+
def read_lfcalc(filepath_or_buffer: Union[str, Path, Any]) -> "CalcSet":
|
250
|
+
"""
|
251
|
+
Read a Leapfrog .lfcalc file and return a CalcSet.
|
252
|
+
|
253
|
+
Args:
|
254
|
+
filepath_or_buffer (str, Path, or file-like): Input file path or buffer.
|
255
|
+
|
256
|
+
Returns:
|
257
|
+
CalcSet: Instance of CalcSet.
|
258
|
+
"""
|
259
|
+
if isinstance(filepath_or_buffer, (str, Path)):
|
260
|
+
with open(filepath_or_buffer, "rb") as f:
|
261
|
+
return CalcSet._read_from_file(f)
|
262
|
+
else:
|
263
|
+
return CalcSet._read_from_file(filepath_or_buffer)
|
264
|
+
|
265
|
+
@staticmethod
|
266
|
+
def _read_from_file(file: Any) -> "CalcSet":
|
267
|
+
"""
|
268
|
+
Read a CalcSet from a file object.
|
269
|
+
Args:
|
270
|
+
file (file-like): File object to read from.
|
271
|
+
Returns:
|
272
|
+
CalcSet: Instance of CalcSet.
|
273
|
+
"""
|
274
|
+
file.seek(len(HEADER))
|
275
|
+
compressed_data = file.read()
|
276
|
+
json_data = zlib.decompress(compressed_data).decode("utf-8")
|
277
|
+
data = json.loads(json_data)
|
278
|
+
return CalcSet.from_dict(data)
|
279
|
+
|
280
|
+
def __repr__(self) -> str:
|
281
|
+
"""
|
282
|
+
Return a pretty-printed JSON string representation of the CalcSet.
|
283
|
+
"""
|
284
|
+
return self.to_json(indent=2)
|
285
|
+
|
286
|
+
def __add__(self, other: "CalcSet") -> "CalcSet":
|
287
|
+
"""
|
288
|
+
Add two CalcSet objects together, combining their items.
|
289
|
+
Args:
|
290
|
+
other (CalcSet): Another CalcSet instance.
|
291
|
+
Returns:
|
292
|
+
CalcSet: New CalcSet with combined items.
|
293
|
+
"""
|
294
|
+
if not isinstance(other, CalcSet):
|
295
|
+
return NotImplemented
|
296
|
+
items1 = list(self.items) if self.items else []
|
297
|
+
items2 = list(other.items) if other.items else []
|
298
|
+
return CalcSet(items1 + items2)
|
299
|
+
|
300
|
+
def rename(
|
301
|
+
self,
|
302
|
+
items: Optional[Union[Dict[str, str], Callable[[str], Optional[str]]]] = None,
|
303
|
+
variables: Optional[
|
304
|
+
Union[Dict[str, str], Callable[[str], Optional[str]]]
|
305
|
+
] = None,
|
306
|
+
regex: bool = False,
|
307
|
+
) -> "CalcSet":
|
308
|
+
"""
|
309
|
+
Return a copy of the CalcSet with specified items renamed and/or variables in children renamed.
|
310
|
+
|
311
|
+
Args:
|
312
|
+
items (dict-like or function): Mapping of old item names to new names.
|
313
|
+
variables (dict-like or function): Mapping of old variable names to new names.
|
314
|
+
regex (bool): Whether to treat keys in `items` and `variables` as regex patterns.
|
315
|
+
|
316
|
+
Returns:
|
317
|
+
CalcSet: New instance with updated item names and/or children.
|
318
|
+
"""
|
319
|
+
new_items = []
|
320
|
+
for item in self.items:
|
321
|
+
name = item.name
|
322
|
+
# Rename item names
|
323
|
+
if items is not None:
|
324
|
+
if callable(items):
|
325
|
+
new_name = items(name)
|
326
|
+
if new_name is not None:
|
327
|
+
name = new_name
|
328
|
+
elif regex:
|
329
|
+
for pattern, replacement in items.items():
|
330
|
+
new_name = re.sub(pattern, replacement, name)
|
331
|
+
if new_name != name:
|
332
|
+
name = new_name
|
333
|
+
else:
|
334
|
+
if name in items:
|
335
|
+
name = items[name]
|
336
|
+
# Rename item name using variables mapping for any Item subclass
|
337
|
+
var_name = name
|
338
|
+
if variables is not None and isinstance(item, Item):
|
339
|
+
if callable(variables):
|
340
|
+
new_var_name = variables(var_name)
|
341
|
+
if new_var_name is not None:
|
342
|
+
var_name = new_var_name
|
343
|
+
elif regex:
|
344
|
+
for pattern, replacement in variables.items():
|
345
|
+
new_var_name = re.sub(pattern, replacement, var_name)
|
346
|
+
if new_var_name != var_name:
|
347
|
+
var_name = new_var_name
|
348
|
+
else:
|
349
|
+
if var_name in variables:
|
350
|
+
var_name = variables[var_name]
|
351
|
+
# Use var_name for all Item subclasses
|
352
|
+
final_name = var_name if isinstance(item, Item) else name
|
353
|
+
new_items.append(
|
354
|
+
item.rename(name=final_name, variables=variables, regex=regex) # type: ignore
|
355
|
+
)
|
356
|
+
return CalcSet(new_items)
|
357
|
+
|
358
|
+
def _repr_html_(self):
|
359
|
+
from .display import display_calcset
|
360
|
+
return display_calcset(self, display_output=False)
|
361
|
+
|
362
|
+
|
363
|
+
class Item:
|
364
|
+
"""
|
365
|
+
Base class for all items in a CalcSet.
|
366
|
+
|
367
|
+
Subclasses represent specific calculation types (Number, Category, Variable, Filter, If, etc.).
|
368
|
+
Each item has a name, a list of child expressions, and optional comments.
|
369
|
+
|
370
|
+
Attributes:
|
371
|
+
name (str): Name of the item.
|
372
|
+
children (list): List of child expressions/statements.
|
373
|
+
comment_item (str): Comment for the item.
|
374
|
+
comment_equation (str): Comment for the equation.
|
375
|
+
"""
|
376
|
+
|
377
|
+
item_type = None
|
378
|
+
calculation_type = None
|
379
|
+
|
380
|
+
def __init__(
|
381
|
+
self,
|
382
|
+
name: str,
|
383
|
+
children: List[Any],
|
384
|
+
comment_item: str = "",
|
385
|
+
comment_equation: str = "",
|
386
|
+
):
|
387
|
+
"""
|
388
|
+
Initialize an Item.
|
389
|
+
Args:
|
390
|
+
name (str): Name of the item.
|
391
|
+
children (list): List of child expressions/statements.
|
392
|
+
comment_item (str): Comment for the item.
|
393
|
+
comment_equation (str): Comment for the equation.
|
394
|
+
"""
|
395
|
+
self.name = name
|
396
|
+
self.children = ensure_list(children)
|
397
|
+
self.comment_item = comment_item
|
398
|
+
self.comment_equation = comment_equation
|
399
|
+
|
400
|
+
def to_dict(self) -> Dict[str, Any]:
|
401
|
+
"""
|
402
|
+
Convert the Item to a dictionary representation.
|
403
|
+
|
404
|
+
Returns:
|
405
|
+
dict: Dictionary representation of the item.
|
406
|
+
"""
|
407
|
+
if self.item_type is None:
|
408
|
+
raise NotImplementedError("item_type must be defined in subclass")
|
409
|
+
children = to_dict(self.children, guard_strings=True)
|
410
|
+
item = {
|
411
|
+
"type": self.item_type,
|
412
|
+
"name": self.name,
|
413
|
+
"equation": {
|
414
|
+
"type": "equation",
|
415
|
+
"comment": self.comment_equation,
|
416
|
+
"statement": {
|
417
|
+
"type": "list",
|
418
|
+
"children": children,
|
419
|
+
},
|
420
|
+
},
|
421
|
+
"comment": self.comment_item,
|
422
|
+
}
|
423
|
+
if self.calculation_type:
|
424
|
+
item["calculation_type"] = self.calculation_type
|
425
|
+
return item
|
426
|
+
|
427
|
+
@classmethod
|
428
|
+
def from_dict(cls: Type["Item"], data: Dict[str, Any]) -> "Item":
|
429
|
+
"""
|
430
|
+
Create an Item from a dictionary.
|
431
|
+
|
432
|
+
Args:
|
433
|
+
data (dict): Dictionary containing item data.
|
434
|
+
|
435
|
+
Returns:
|
436
|
+
Item: Instance of Item or subclass.
|
437
|
+
"""
|
438
|
+
if cls.item_type is None:
|
439
|
+
raise NotImplementedError("item_type must be defined in subclass")
|
440
|
+
if data["type"] != cls.item_type:
|
441
|
+
raise ValueError(f"Expected item type {cls.item_type}, got {data['type']}")
|
442
|
+
children = []
|
443
|
+
for child in ensure_list(data["equation"]["statement"]["children"]):
|
444
|
+
children.append(dispatch_expression(child))
|
445
|
+
return cls(
|
446
|
+
name=data["name"],
|
447
|
+
children=children,
|
448
|
+
comment_item=data.get("comment", ""),
|
449
|
+
comment_equation=data["equation"].get("comment", ""),
|
450
|
+
)
|
451
|
+
|
452
|
+
@property
|
453
|
+
def dependencies(self) -> Set[str]:
|
454
|
+
"""
|
455
|
+
Get the set of variable dependencies for this item.
|
456
|
+
|
457
|
+
Returns:
|
458
|
+
set: Set of variable names that are dependencies.
|
459
|
+
"""
|
460
|
+
return get_dependencies(self)
|
461
|
+
|
462
|
+
def copy(self) -> "Item":
|
463
|
+
"""
|
464
|
+
Return a deep copy of the Item.
|
465
|
+
"""
|
466
|
+
return type(self)(
|
467
|
+
name=self.name,
|
468
|
+
children=[c.copy() if hasattr(c, "copy") else c for c in self.children],
|
469
|
+
comment_item=self.comment_item,
|
470
|
+
comment_equation=self.comment_equation,
|
471
|
+
)
|
472
|
+
|
473
|
+
def replace(self, **changes: Any) -> "Item":
|
474
|
+
"""
|
475
|
+
Return a copy of the Item with specified attributes replaced.
|
476
|
+
|
477
|
+
Args:
|
478
|
+
**changes: Attributes to replace.
|
479
|
+
|
480
|
+
Returns:
|
481
|
+
Item: New instance with updated attributes.
|
482
|
+
"""
|
483
|
+
params = {
|
484
|
+
"name": self.name,
|
485
|
+
"children": self.children,
|
486
|
+
"comment_item": self.comment_item,
|
487
|
+
"comment_equation": self.comment_equation,
|
488
|
+
}
|
489
|
+
params.update(changes)
|
490
|
+
return type(self)(**params)
|
491
|
+
|
492
|
+
def rename(
|
493
|
+
self,
|
494
|
+
name: Optional[str] = None,
|
495
|
+
variables: Optional[
|
496
|
+
Union[Dict[str, str], Callable[[str], Optional[str]]]
|
497
|
+
] = None,
|
498
|
+
regex: bool = False,
|
499
|
+
) -> "Item":
|
500
|
+
"""
|
501
|
+
Return a copy of the Item with a new name and/or renamed variables in children.
|
502
|
+
|
503
|
+
Args:
|
504
|
+
name (str): New name for the item.
|
505
|
+
variables (dict-like or function): Mapping of old variable names to new names.
|
506
|
+
regex (bool): Whether to treat keys in `variables` as regex patterns.
|
507
|
+
|
508
|
+
Returns:
|
509
|
+
Item: New instance with updated name and/or children.
|
510
|
+
"""
|
511
|
+
new = self.copy()
|
512
|
+
# For any Item subclass, allow variable renaming to affect the name
|
513
|
+
if name is not None:
|
514
|
+
new.name = name
|
515
|
+
elif variables is not None and isinstance(self, Item):
|
516
|
+
var_name = new.name
|
517
|
+
if callable(variables):
|
518
|
+
new_var_name = variables(var_name)
|
519
|
+
if new_var_name is not None:
|
520
|
+
new.name = new_var_name
|
521
|
+
elif regex:
|
522
|
+
for pattern, replacement in variables.items():
|
523
|
+
new_var_name = re.sub(pattern, replacement, var_name)
|
524
|
+
if new_var_name != var_name:
|
525
|
+
var_name = new_var_name
|
526
|
+
new.name = var_name
|
527
|
+
else:
|
528
|
+
if var_name in variables:
|
529
|
+
new.name = variables[var_name]
|
530
|
+
if variables is not None:
|
531
|
+
return rename(new, variables, regex=regex) # type: ignore
|
532
|
+
return new
|
533
|
+
|
534
|
+
|
535
|
+
class IfRow:
|
536
|
+
"""
|
537
|
+
Represents a single row in an If block, containing a condition and corresponding value(s).
|
538
|
+
|
539
|
+
Args:
|
540
|
+
condition (list): Condition expressions.
|
541
|
+
value (list): Value expressions if condition is met.
|
542
|
+
"""
|
543
|
+
def __init__(self, condition: List[Any], value: List[Any]):
|
544
|
+
"""
|
545
|
+
Initialize an IfRow.
|
546
|
+
Args:
|
547
|
+
condition (list): Condition expressions.
|
548
|
+
value (list): Value expressions if condition is met.
|
549
|
+
"""
|
550
|
+
self.condition = condition
|
551
|
+
self.value = value
|
552
|
+
|
553
|
+
def to_dict(self) -> Dict[str, Any]:
|
554
|
+
"""
|
555
|
+
Convert the IfRow to a dictionary representation.
|
556
|
+
|
557
|
+
Returns:
|
558
|
+
dict: Dictionary representation of the IfRow.
|
559
|
+
"""
|
560
|
+
return {
|
561
|
+
"type": "if_row",
|
562
|
+
"test": {"type": "list", "children": to_dict(self.condition)},
|
563
|
+
"result": {
|
564
|
+
"type": "list",
|
565
|
+
"children": to_dict(self.value, guard_strings=True),
|
566
|
+
},
|
567
|
+
}
|
568
|
+
|
569
|
+
@classmethod
|
570
|
+
def from_dict(cls: Type["IfRow"], data: Dict[str, Any]) -> "IfRow":
|
571
|
+
"""
|
572
|
+
Create an IfRow from a dictionary.
|
573
|
+
|
574
|
+
Args:
|
575
|
+
data (dict): Dictionary containing IfRow data.
|
576
|
+
|
577
|
+
Returns:
|
578
|
+
IfRow: Instance of IfRow.
|
579
|
+
"""
|
580
|
+
if data["type"] != "if_row":
|
581
|
+
raise ValueError(f"Expected type 'if_row', got {data['type']}")
|
582
|
+
# return cls(
|
583
|
+
# condition=data["test"]["children"],
|
584
|
+
# value=data["result"]["children"],
|
585
|
+
# )
|
586
|
+
condition = []
|
587
|
+
for cond in ensure_list(data["test"]["children"]):
|
588
|
+
condition.append(dispatch_expression(cond))
|
589
|
+
value = []
|
590
|
+
for val in ensure_str_list(data["result"]["children"]):
|
591
|
+
value.append(dispatch_expression(val))
|
592
|
+
return cls(condition=condition, value=value)
|
593
|
+
|
594
|
+
def copy(self) -> "IfRow":
|
595
|
+
"""
|
596
|
+
Return a deep copy of the IfRow.
|
597
|
+
"""
|
598
|
+
return IfRow(
|
599
|
+
condition=[c.copy() if hasattr(c, "copy") else c for c in self.condition],
|
600
|
+
value=[v.copy() if hasattr(v, "copy") else v for v in self.value],
|
601
|
+
)
|
602
|
+
|
603
|
+
|
604
|
+
class If:
|
605
|
+
"""
|
606
|
+
Represents a conditional logic block (if/else) in a calculation set.
|
607
|
+
|
608
|
+
Args:
|
609
|
+
rows (list): List of IfRow objects, dicts, or (condition, value) tuples.
|
610
|
+
otherwise (list): Expressions for the 'otherwise' case.
|
611
|
+
|
612
|
+
Example:
|
613
|
+
>>> from pollywog.core import If, Number
|
614
|
+
>>> if_block = If([
|
615
|
+
... (["[Au] > 1"], "[Au] * 1.1"),
|
616
|
+
... (["[Au] <= 1"], "[Au] * 0.9")
|
617
|
+
... ], otherwise=["[Au]"])
|
618
|
+
"""
|
619
|
+
def __init__(self, rows: List[Any], otherwise: List[Any]):
|
620
|
+
"""
|
621
|
+
Initialize an If expression.
|
622
|
+
Args:
|
623
|
+
rows (list): List of either IfRow objects, dicts, or (condition, value) tuples.
|
624
|
+
otherwise (list): Expressions for the 'otherwise' case.
|
625
|
+
"""
|
626
|
+
self.rows = rows
|
627
|
+
self.otherwise = otherwise
|
628
|
+
|
629
|
+
def to_dict(self) -> Dict[str, Any]:
|
630
|
+
"""
|
631
|
+
Convert the If expression to a dictionary representation.
|
632
|
+
|
633
|
+
Returns:
|
634
|
+
dict: Dictionary representation of the If expression.
|
635
|
+
"""
|
636
|
+
rows = []
|
637
|
+
for row in ensure_list(self.rows):
|
638
|
+
if isinstance(row, IfRow):
|
639
|
+
rows.append(row.to_dict())
|
640
|
+
elif isinstance(row, dict) and row.get("type") == "if_row":
|
641
|
+
rows.append(row)
|
642
|
+
elif isinstance(row, (tuple, list)) and len(row) == 2:
|
643
|
+
condition, value = row
|
644
|
+
rows.append(IfRow(condition, value).to_dict())
|
645
|
+
else:
|
646
|
+
raise ValueError(f"Invalid row format: {row}")
|
647
|
+
return {
|
648
|
+
"type": "if",
|
649
|
+
"rows": rows,
|
650
|
+
"otherwise": {
|
651
|
+
"type": "list",
|
652
|
+
"children": to_dict(self.otherwise, guard_strings=True),
|
653
|
+
},
|
654
|
+
}
|
655
|
+
|
656
|
+
@classmethod
|
657
|
+
def from_dict(cls: Type["If"], data: Dict[str, Any]) -> "If":
|
658
|
+
"""
|
659
|
+
Create an If expression from a dictionary.
|
660
|
+
|
661
|
+
Args:
|
662
|
+
data (dict): Dictionary containing If expression data.
|
663
|
+
|
664
|
+
Returns:
|
665
|
+
If: Instance of If.
|
666
|
+
"""
|
667
|
+
if data["type"] != "if":
|
668
|
+
raise ValueError(f"Expected type 'if', got {data['type']}")
|
669
|
+
rows = [
|
670
|
+
IfRow.from_dict(row) if isinstance(row, dict) else row
|
671
|
+
for row in ensure_list(data["rows"])
|
672
|
+
]
|
673
|
+
otherwise = []
|
674
|
+
for val in ensure_str_list(data["otherwise"]["children"]):
|
675
|
+
otherwise.append(dispatch_expression(val))
|
676
|
+
return cls(rows=rows, otherwise=otherwise)
|
677
|
+
|
678
|
+
def copy(self) -> "If":
|
679
|
+
"""
|
680
|
+
Return a deep copy of the If expression.
|
681
|
+
"""
|
682
|
+
return If(
|
683
|
+
rows=[r.copy() if hasattr(r, "copy") else r for r in self.rows],
|
684
|
+
otherwise=[o.copy() if hasattr(o, "copy") else o for o in self.otherwise],
|
685
|
+
)
|
686
|
+
|
687
|
+
|
688
|
+
class Number(Item):
|
689
|
+
"""
|
690
|
+
Represents a numeric calculation item in a CalcSet.
|
691
|
+
|
692
|
+
Used for variables whose values are numbers, either integers or floats.
|
693
|
+
"""
|
694
|
+
|
695
|
+
item_type = "calculation"
|
696
|
+
calculation_type = "number"
|
697
|
+
|
698
|
+
|
699
|
+
class Category(Item):
|
700
|
+
"""
|
701
|
+
Represents a categorical calculation item in a CalcSet.
|
702
|
+
|
703
|
+
Used for variables whose values are categories, represented as strings.
|
704
|
+
"""
|
705
|
+
|
706
|
+
item_type = "calculation"
|
707
|
+
calculation_type = "string"
|
708
|
+
|
709
|
+
|
710
|
+
class Variable(Item):
|
711
|
+
"""
|
712
|
+
Represents a variable item in a calculation set.
|
713
|
+
|
714
|
+
Used for declaring variables that may be referenced by other calculations.
|
715
|
+
"""
|
716
|
+
|
717
|
+
item_type = "variable"
|
718
|
+
|
719
|
+
|
720
|
+
class Filter(Item):
|
721
|
+
"""
|
722
|
+
Represents a filter item in a calculation set.
|
723
|
+
|
724
|
+
Used for defining filters that restrict or modify calculation results.
|
725
|
+
"""
|
726
|
+
|
727
|
+
item_type = "filter"
|
728
|
+
|
729
|
+
|
730
|
+
classes = {
|
731
|
+
# "calculation": Item,
|
732
|
+
"variable": Variable,
|
733
|
+
"filter": Filter,
|
734
|
+
"if": If,
|
735
|
+
"if_row": IfRow,
|
736
|
+
}
|
737
|
+
|
738
|
+
|
739
|
+
expressions = {
|
740
|
+
"if": If,
|
741
|
+
}
|
742
|
+
|
743
|
+
|
744
|
+
def dispatch_expression(data: Any) -> Any:
|
745
|
+
"""
|
746
|
+
Dispatch an expression dictionary to the appropriate class constructor.
|
747
|
+
|
748
|
+
Args:
|
749
|
+
data (dict or any): Expression data.
|
750
|
+
|
751
|
+
Returns:
|
752
|
+
object: Instantiated expression object or the original data if not a dict.
|
753
|
+
"""
|
754
|
+
if isinstance(data, dict) and "type" in data:
|
755
|
+
expr_type = data["type"]
|
756
|
+
if expr_type in expressions:
|
757
|
+
return expressions[expr_type].from_dict(data)
|
758
|
+
else:
|
759
|
+
raise ValueError(f"Unknown expression type: {expr_type}")
|
760
|
+
return data
|
761
|
+
|
762
|
+
|
763
|
+
def get_dependencies(item: Any) -> Set[str]:
|
764
|
+
"""
|
765
|
+
Recursively extract variable dependencies from an Item or expression.
|
766
|
+
|
767
|
+
Args:
|
768
|
+
item (Item or expression): The item or expression to analyze.
|
769
|
+
|
770
|
+
Returns:
|
771
|
+
set: Set of variable names that are dependencies.
|
772
|
+
"""
|
773
|
+
deps = set()
|
774
|
+
|
775
|
+
if isinstance(item, Item):
|
776
|
+
for child in item.children:
|
777
|
+
deps.update(get_dependencies(child))
|
778
|
+
elif isinstance(item, If):
|
779
|
+
for row in item.rows:
|
780
|
+
deps.update(get_dependencies(row))
|
781
|
+
deps.update(get_dependencies(item.otherwise))
|
782
|
+
elif isinstance(item, IfRow):
|
783
|
+
deps.update(get_dependencies(item.condition))
|
784
|
+
deps.update(get_dependencies(item.value))
|
785
|
+
elif isinstance(item, list):
|
786
|
+
for elem in item:
|
787
|
+
deps.update(get_dependencies(elem))
|
788
|
+
elif isinstance(item, str):
|
789
|
+
# Find all occurrences of [var_name] in the string
|
790
|
+
found_vars = re.findall(r"\[([^\[\]]+)\]", item)
|
791
|
+
deps.update(found_vars)
|
792
|
+
|
793
|
+
return deps
|
794
|
+
|
795
|
+
|
796
|
+
def rename(
|
797
|
+
item: Any, mapper: Union[Dict[str, str], Callable[[str], str]], regex: bool = False
|
798
|
+
) -> Any:
|
799
|
+
"""
|
800
|
+
Recursively rename variables in an Item or expression based on a mapping dictionary.
|
801
|
+
|
802
|
+
Args:
|
803
|
+
item (Item or expression): The item or expression to rename.
|
804
|
+
mapper (dict-like or function): Mapping of old variable names to new names.
|
805
|
+
regex (bool): If True, treat keys and values of the mapper as regular expressions.
|
806
|
+
|
807
|
+
Returns:
|
808
|
+
Item or expression: The renamed item or expression.
|
809
|
+
"""
|
810
|
+
if isinstance(item, Item):
|
811
|
+
new_children = [rename(child, mapper, regex=regex) for child in item.children]
|
812
|
+
return item.replace(children=new_children)
|
813
|
+
elif isinstance(item, If):
|
814
|
+
new_rows = [rename(row, mapper, regex=regex) for row in item.rows]
|
815
|
+
new_otherwise = rename(item.otherwise, mapper, regex=regex)
|
816
|
+
return If(rows=new_rows, otherwise=new_otherwise)
|
817
|
+
elif isinstance(item, IfRow):
|
818
|
+
new_condition = rename(item.condition, mapper, regex=regex)
|
819
|
+
new_value = rename(item.value, mapper, regex=regex)
|
820
|
+
return IfRow(condition=new_condition, value=new_value)
|
821
|
+
elif isinstance(item, list):
|
822
|
+
return [rename(elem, mapper, regex=regex) for elem in item]
|
823
|
+
elif isinstance(item, str):
|
824
|
+
|
825
|
+
def replace_var(match):
|
826
|
+
var_name = match.group(1)
|
827
|
+
if callable(mapper):
|
828
|
+
return f"[{mapper(var_name)}]"
|
829
|
+
elif regex:
|
830
|
+
for pattern, replacement in mapper.items():
|
831
|
+
var_name = re.sub(pattern, replacement, var_name)
|
832
|
+
return f"[{var_name}]"
|
833
|
+
else:
|
834
|
+
return f"[{mapper.get(var_name, var_name)}]"
|
835
|
+
|
836
|
+
return re.sub(r"\[([^\[\]]+)\]", replace_var, item)
|
837
|
+
else:
|
838
|
+
return item
|