codestr 0.3.0__tar.gz → 0.3.2__tar.gz
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.
- {codestr-0.3.0 → codestr-0.3.2}/.github/workflows/ci.yml +17 -0
- {codestr-0.3.0 → codestr-0.3.2}/PKG-INFO +23 -4
- {codestr-0.3.0 → codestr-0.3.2}/README.md +20 -1
- {codestr-0.3.0 → codestr-0.3.2}/pyproject.toml +2 -2
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/engine.py +118 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/udf/cs_udf.py +1 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_base_udf.py +98 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_cs_udf.py +41 -1
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_engine.py +127 -0
- {codestr-0.3.0 → codestr-0.3.2}/uv.lock +15 -15
- {codestr-0.3.0 → codestr-0.3.2}/.github/workflows/release.yml +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/.gitignore +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/.python-version +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/CHANGELOG.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/CLAUDE.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/CONTRIBUTING.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/LICENSE +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/operators.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/plans/2026-07-16-ts-ema-keyword-arguments.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/plans/2026-07-17-min-samples.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/plans/2026-07-17-mixed-window-planner.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-16-polars-minimum-version-design.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-16-ts-ema-keyword-arguments-design.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-17-min-samples-design.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-17-mixed-window-planner-design.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/review.md +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/__init__.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/compiler.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/errors.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/parser.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/planner.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/syntax.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/tokens.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/udf/__init__.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/udf/base_udf.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/udf/registry.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/src/codestr/udf/ts_udf.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/__init__.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/conftest.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_compiler.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_parser.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_planner.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_registry.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_syntax.py +0 -0
- {codestr-0.3.0 → codestr-0.3.2}/tests/test_ts_udf.py +0 -0
|
@@ -33,3 +33,20 @@ jobs:
|
|
|
33
33
|
run: |
|
|
34
34
|
uv sync --extra test
|
|
35
35
|
uv run pytest tests/ -v
|
|
36
|
+
|
|
37
|
+
coverage:
|
|
38
|
+
runs-on: ubuntu-latest
|
|
39
|
+
steps:
|
|
40
|
+
- uses: actions/checkout@v4
|
|
41
|
+
- uses: astral-sh/setup-uv@v5
|
|
42
|
+
with:
|
|
43
|
+
python-version: "3.12"
|
|
44
|
+
- name: Install test dependencies
|
|
45
|
+
run: uv sync --extra test
|
|
46
|
+
- name: Enforce branch coverage
|
|
47
|
+
run: |
|
|
48
|
+
uv run pytest tests/ \
|
|
49
|
+
--cov=src/codestr \
|
|
50
|
+
--cov-branch \
|
|
51
|
+
--cov-report=term-missing \
|
|
52
|
+
--cov-fail-under=90
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: codestr
|
|
3
|
-
Version: 0.3.
|
|
3
|
+
Version: 0.3.2
|
|
4
4
|
Summary: DSL → Polars Expr compile engine for quantitative factor mining
|
|
5
5
|
Project-URL: Repository, https://github.com/huangbogeng/codestr
|
|
6
6
|
Author: huangbogeng
|
|
@@ -18,7 +18,7 @@ Requires-Python: >=3.10
|
|
|
18
18
|
Requires-Dist: lark<2.0,>=1.0
|
|
19
19
|
Requires-Dist: loguru<1.0,>=0.7
|
|
20
20
|
Requires-Dist: numpy<3.0,>=1.26
|
|
21
|
-
Requires-Dist: polars>=1.
|
|
21
|
+
Requires-Dist: polars>=1.44.1
|
|
22
22
|
Requires-Dist: toolz<2.0,>=1.0
|
|
23
23
|
Provides-Extra: dev
|
|
24
24
|
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
@@ -71,12 +71,13 @@ result = cs.sql(
|
|
|
71
71
|
print(result)
|
|
72
72
|
```
|
|
73
73
|
|
|
74
|
-
##
|
|
74
|
+
## API 模式
|
|
75
75
|
|
|
76
76
|
| 模式 | API | 行为 |
|
|
77
77
|
|------|-----|------|
|
|
78
78
|
| **纯编译** | `cs.compile(expr) -> pl.Expr` | 无副作用,返回 Polars 表达式 |
|
|
79
79
|
| **交互式** | `cs.sql(expr, lazy=False) -> pl.DataFrame` | 有状态,自动缓存与复用 |
|
|
80
|
+
| **静态验证** | `cs.validate_expr(*exprs) -> list[dict]` | 无副作用,按当前 schema 干编译 |
|
|
80
81
|
|
|
81
82
|
```python
|
|
82
83
|
# 纯编译 — 表达式可被任意 DataFrame 消费
|
|
@@ -88,6 +89,24 @@ cs.sql("close + volume as total")
|
|
|
88
89
|
cs.sql("ts_mean(total, 5) as total_ma5") # 复用上一步的 total
|
|
89
90
|
```
|
|
90
91
|
|
|
92
|
+
### 静态验证
|
|
93
|
+
|
|
94
|
+
`validate_expr()` 复用 `sql()` 的 planner,在当前 LazyFrame schema 上解析
|
|
95
|
+
UDF、列引用、类型兼容性和混合窗口,但只调用 `collect_schema()`,不会
|
|
96
|
+
`collect()` 数据或修改引擎缓存:
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
results = cs.validate_expr(
|
|
100
|
+
"sin(1.0) as invalid",
|
|
101
|
+
"ts_mean(cs_moderate(close), 5) as factor",
|
|
102
|
+
)
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
每个结果包含 `expr`、`valid`、`stage`、`error_type` 和 `message`。
|
|
106
|
+
失败阶段分为 `structural`、`compile` 和 `schema`。批量输入相互独立,
|
|
107
|
+
后一条表达式不能引用同批前一条表达式新建的别名。依赖实际数据值、
|
|
108
|
+
只在物化时出现的错误不属于该 API 的检查范围。
|
|
109
|
+
|
|
91
110
|
## 窗口配置
|
|
92
111
|
|
|
93
112
|
CodeStr 使用 `partition_by`(实体分组轴)和 `order_by`(时间排序轴)控制窗口算子:
|
|
@@ -40,12 +40,13 @@ result = cs.sql(
|
|
|
40
40
|
print(result)
|
|
41
41
|
```
|
|
42
42
|
|
|
43
|
-
##
|
|
43
|
+
## API 模式
|
|
44
44
|
|
|
45
45
|
| 模式 | API | 行为 |
|
|
46
46
|
|------|-----|------|
|
|
47
47
|
| **纯编译** | `cs.compile(expr) -> pl.Expr` | 无副作用,返回 Polars 表达式 |
|
|
48
48
|
| **交互式** | `cs.sql(expr, lazy=False) -> pl.DataFrame` | 有状态,自动缓存与复用 |
|
|
49
|
+
| **静态验证** | `cs.validate_expr(*exprs) -> list[dict]` | 无副作用,按当前 schema 干编译 |
|
|
49
50
|
|
|
50
51
|
```python
|
|
51
52
|
# 纯编译 — 表达式可被任意 DataFrame 消费
|
|
@@ -57,6 +58,24 @@ cs.sql("close + volume as total")
|
|
|
57
58
|
cs.sql("ts_mean(total, 5) as total_ma5") # 复用上一步的 total
|
|
58
59
|
```
|
|
59
60
|
|
|
61
|
+
### 静态验证
|
|
62
|
+
|
|
63
|
+
`validate_expr()` 复用 `sql()` 的 planner,在当前 LazyFrame schema 上解析
|
|
64
|
+
UDF、列引用、类型兼容性和混合窗口,但只调用 `collect_schema()`,不会
|
|
65
|
+
`collect()` 数据或修改引擎缓存:
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
results = cs.validate_expr(
|
|
69
|
+
"sin(1.0) as invalid",
|
|
70
|
+
"ts_mean(cs_moderate(close), 5) as factor",
|
|
71
|
+
)
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
每个结果包含 `expr`、`valid`、`stage`、`error_type` 和 `message`。
|
|
75
|
+
失败阶段分为 `structural`、`compile` 和 `schema`。批量输入相互独立,
|
|
76
|
+
后一条表达式不能引用同批前一条表达式新建的别名。依赖实际数据值、
|
|
77
|
+
只在物化时出现的错误不属于该 API 的检查范围。
|
|
78
|
+
|
|
60
79
|
## 窗口配置
|
|
61
80
|
|
|
62
81
|
CodeStr 使用 `partition_by`(实体分组轴)和 `order_by`(时间排序轴)控制窗口算子:
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "codestr"
|
|
7
|
-
version = "0.3.
|
|
7
|
+
version = "0.3.2"
|
|
8
8
|
description = "DSL → Polars Expr compile engine for quantitative factor mining"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
license = {text = "MIT"}
|
|
@@ -23,7 +23,7 @@ classifiers = [
|
|
|
23
23
|
"Programming Language :: Python :: 3.12",
|
|
24
24
|
]
|
|
25
25
|
dependencies = [
|
|
26
|
-
"polars>=1.
|
|
26
|
+
"polars>=1.44.1",
|
|
27
27
|
"lark>=1.0,<2.0",
|
|
28
28
|
"loguru>=0.7,<1.0",
|
|
29
29
|
"toolz>=1.0,<2.0",
|
|
@@ -195,6 +195,124 @@ class CodeStr:
|
|
|
195
195
|
result["valid"] = len(result["reasons"]) == 0
|
|
196
196
|
return result
|
|
197
197
|
|
|
198
|
+
def validate_expr(self, *exprs: str) -> list[dict[str, object]]:
|
|
199
|
+
"""Dry-compile expressions against the current Polars schema.
|
|
200
|
+
|
|
201
|
+
Unlike :meth:`check_expr`, this resolves UDFs, planner stages, and
|
|
202
|
+
column/type compatibility. It never collects data or mutates engine
|
|
203
|
+
state. Batch expressions are validated independently.
|
|
204
|
+
|
|
205
|
+
Returns:
|
|
206
|
+
One result per input expression. ``stage`` is ``structural``,
|
|
207
|
+
``compile``, or ``schema``; successful validation reaches the
|
|
208
|
+
``schema`` stage with no error details.
|
|
209
|
+
"""
|
|
210
|
+
base = self._data_ if self._data_ is not None else self.data.lazy()
|
|
211
|
+
registry = UDFRegistry.get_instance()
|
|
212
|
+
results = []
|
|
213
|
+
|
|
214
|
+
for expr in exprs:
|
|
215
|
+
try:
|
|
216
|
+
node = _parse(expr)
|
|
217
|
+
except Exception as error:
|
|
218
|
+
results.append(
|
|
219
|
+
{
|
|
220
|
+
"expr": expr,
|
|
221
|
+
"valid": False,
|
|
222
|
+
"stage": "structural",
|
|
223
|
+
"error_type": type(error).__name__,
|
|
224
|
+
"message": str(error),
|
|
225
|
+
}
|
|
226
|
+
)
|
|
227
|
+
continue
|
|
228
|
+
|
|
229
|
+
structural = self.check_expr(expr)
|
|
230
|
+
if not structural["valid"]:
|
|
231
|
+
results.append(
|
|
232
|
+
{
|
|
233
|
+
"expr": expr,
|
|
234
|
+
"valid": False,
|
|
235
|
+
"stage": "structural",
|
|
236
|
+
"error_type": "StructuralError",
|
|
237
|
+
"message": "; ".join(structural["reasons"]),
|
|
238
|
+
}
|
|
239
|
+
)
|
|
240
|
+
continue
|
|
241
|
+
|
|
242
|
+
try:
|
|
243
|
+
current_columns = base.collect_schema().names()
|
|
244
|
+
except Exception as error:
|
|
245
|
+
results.append(
|
|
246
|
+
{
|
|
247
|
+
"expr": expr,
|
|
248
|
+
"valid": False,
|
|
249
|
+
"stage": "schema",
|
|
250
|
+
"error_type": type(error).__name__,
|
|
251
|
+
"message": str(error),
|
|
252
|
+
}
|
|
253
|
+
)
|
|
254
|
+
continue
|
|
255
|
+
|
|
256
|
+
try:
|
|
257
|
+
plan = build_execution_plan(
|
|
258
|
+
node,
|
|
259
|
+
registry,
|
|
260
|
+
existing_columns=current_columns,
|
|
261
|
+
)
|
|
262
|
+
compiled_steps = [
|
|
263
|
+
(
|
|
264
|
+
_pure_compile(
|
|
265
|
+
step.node,
|
|
266
|
+
registry=registry,
|
|
267
|
+
dims=getattr(self, "dims", None),
|
|
268
|
+
ts_over=self._ts_over,
|
|
269
|
+
cs_over=self._cs_over,
|
|
270
|
+
),
|
|
271
|
+
step.output_name,
|
|
272
|
+
)
|
|
273
|
+
for step in plan.steps
|
|
274
|
+
]
|
|
275
|
+
except Exception as error:
|
|
276
|
+
results.append(
|
|
277
|
+
{
|
|
278
|
+
"expr": expr,
|
|
279
|
+
"valid": False,
|
|
280
|
+
"stage": "compile",
|
|
281
|
+
"error_type": type(error).__name__,
|
|
282
|
+
"message": str(error),
|
|
283
|
+
}
|
|
284
|
+
)
|
|
285
|
+
continue
|
|
286
|
+
|
|
287
|
+
try:
|
|
288
|
+
candidate = base
|
|
289
|
+
for compiled, output_name in compiled_steps:
|
|
290
|
+
candidate = candidate.with_columns(compiled.alias(output_name))
|
|
291
|
+
candidate.collect_schema()
|
|
292
|
+
except Exception as error:
|
|
293
|
+
results.append(
|
|
294
|
+
{
|
|
295
|
+
"expr": expr,
|
|
296
|
+
"valid": False,
|
|
297
|
+
"stage": "schema",
|
|
298
|
+
"error_type": type(error).__name__,
|
|
299
|
+
"message": str(error),
|
|
300
|
+
}
|
|
301
|
+
)
|
|
302
|
+
continue
|
|
303
|
+
|
|
304
|
+
results.append(
|
|
305
|
+
{
|
|
306
|
+
"expr": expr,
|
|
307
|
+
"valid": True,
|
|
308
|
+
"stage": "schema",
|
|
309
|
+
"error_type": None,
|
|
310
|
+
"message": None,
|
|
311
|
+
}
|
|
312
|
+
)
|
|
313
|
+
|
|
314
|
+
return results
|
|
315
|
+
|
|
198
316
|
def _check_redundant(self, node: Call, reasons: list[str]):
|
|
199
317
|
if not isinstance(node, Call):
|
|
200
318
|
return
|
|
@@ -73,6 +73,7 @@ def cs_qcut(expr: pl.Expr, n_bins=10, partition_by=None, order_by=None):
|
|
|
73
73
|
return (
|
|
74
74
|
expr.qcut(n_bins, labels=[str(i) for i in range(1, n_bins + 1)], allow_duplicates=True)
|
|
75
75
|
.over(partition_by=partition_by, order_by=order_by)
|
|
76
|
+
.cast(pl.String)
|
|
76
77
|
.cast(pl.Int32)
|
|
77
78
|
)
|
|
78
79
|
|
|
@@ -6,6 +6,7 @@ import pytest
|
|
|
6
6
|
|
|
7
7
|
from codestr.compiler import compile as ast_compile
|
|
8
8
|
from codestr.syntax import Call, Column, Literal
|
|
9
|
+
from codestr.udf import base_udf
|
|
9
10
|
|
|
10
11
|
|
|
11
12
|
def _eval_expr(node, df: pl.DataFrame) -> pl.DataFrame:
|
|
@@ -84,6 +85,75 @@ class TestUnaryMath:
|
|
|
84
85
|
result = df.select(ast_compile(node))["between(x, 2, 8)"].to_list()
|
|
85
86
|
assert result == [False, True, False]
|
|
86
87
|
|
|
88
|
+
@pytest.mark.parametrize(
|
|
89
|
+
("name", "value", "expected"),
|
|
90
|
+
[
|
|
91
|
+
("cbrt", 8.0, 2.0),
|
|
92
|
+
("sinh", 0.0, 0.0),
|
|
93
|
+
("arcsin", 0.0, 0.0),
|
|
94
|
+
("arcsinh", 0.0, 0.0),
|
|
95
|
+
("cosh", 0.0, 1.0),
|
|
96
|
+
("arccos", 1.0, 0.0),
|
|
97
|
+
("arccosh", 1.0, 0.0),
|
|
98
|
+
("tan", 0.0, 0.0),
|
|
99
|
+
("tanh", 0.0, 0.0),
|
|
100
|
+
("arctan", 0.0, 0.0),
|
|
101
|
+
("arctanh", 0.0, 0.0),
|
|
102
|
+
("cot", np.pi / 4, 1.0),
|
|
103
|
+
("degrees", np.pi, 180.0),
|
|
104
|
+
("log1p", np.e - 1, 1.0),
|
|
105
|
+
],
|
|
106
|
+
)
|
|
107
|
+
def test_remaining_unary_numeric_contracts(self, name, value, expected):
|
|
108
|
+
df = pl.DataFrame({"x": [value]})
|
|
109
|
+
result = df.select(ast_compile(Call(name, (Column("x"),))))
|
|
110
|
+
|
|
111
|
+
assert result.item() == pytest.approx(expected)
|
|
112
|
+
|
|
113
|
+
def test_not(self):
|
|
114
|
+
df = pl.DataFrame({"x": [True, False]})
|
|
115
|
+
|
|
116
|
+
result = df.select(ast_compile(Call("not_", (Column("x"),))))
|
|
117
|
+
|
|
118
|
+
assert result.to_series().to_list() == [False, True]
|
|
119
|
+
|
|
120
|
+
def test_entropy(self):
|
|
121
|
+
df = pl.DataFrame({"x": [0.25, 0.75]})
|
|
122
|
+
|
|
123
|
+
result = df.select(ast_compile(Call("entropy", (Column("x"),))))
|
|
124
|
+
|
|
125
|
+
assert result.item() == pytest.approx(0.5623351446)
|
|
126
|
+
|
|
127
|
+
def test_trunc_closed_and_open_bounds(self):
|
|
128
|
+
df = pl.DataFrame({"x": [0.0, 1.0, 2.0]})
|
|
129
|
+
closed = base_udf.trunc(pl.col("x"), 0, 2)
|
|
130
|
+
opened = base_udf.trunc(pl.col("x"), 0, 2, left_closed=False, right_closed=False)
|
|
131
|
+
|
|
132
|
+
result = df.select(closed.alias("closed"), opened.alias("opened"))
|
|
133
|
+
|
|
134
|
+
assert result["closed"].to_list() == [0.0, 1.0, 2.0]
|
|
135
|
+
assert result["opened"].to_list() == [None, 1.0, None]
|
|
136
|
+
|
|
137
|
+
def test_cast_valid_and_invalid_dtype(self):
|
|
138
|
+
df = pl.DataFrame({"x": [1.5]})
|
|
139
|
+
|
|
140
|
+
result = df.select(base_udf.cast(pl.col("x"), "int"))
|
|
141
|
+
|
|
142
|
+
assert result.schema["x"] == pl.Int64
|
|
143
|
+
with pytest.raises(ValueError, match="not a valid type"):
|
|
144
|
+
base_udf.cast(pl.col("x"), "date")
|
|
145
|
+
|
|
146
|
+
def test_concat_and_null_type(self):
|
|
147
|
+
df = pl.DataFrame({"a": [1], "b": [2]})
|
|
148
|
+
|
|
149
|
+
result = df.select(
|
|
150
|
+
base_udf.concat(pl.col("a"), pl.col("b")).alias("values"),
|
|
151
|
+
base_udf.null_type(pl.col("a")).alias("null"),
|
|
152
|
+
)
|
|
153
|
+
|
|
154
|
+
assert result["values"].to_list() == [[1, 2]]
|
|
155
|
+
assert result["null"].to_list() == [None]
|
|
156
|
+
|
|
87
157
|
|
|
88
158
|
class TestBinaryMath:
|
|
89
159
|
def test_add(self):
|
|
@@ -116,6 +186,23 @@ class TestBinaryMath:
|
|
|
116
186
|
node = Call("eq", (Column("a"), Column("b")))
|
|
117
187
|
assert df.select(ast_compile(node))["(a==b)"].to_list() == [True, False]
|
|
118
188
|
|
|
189
|
+
@pytest.mark.parametrize(
|
|
190
|
+
("name", "expected"),
|
|
191
|
+
[
|
|
192
|
+
("floordiv", [2, 2]),
|
|
193
|
+
("mod", [1, 1]),
|
|
194
|
+
("lt", [False, False]),
|
|
195
|
+
("le", [False, False]),
|
|
196
|
+
("ge", [True, True]),
|
|
197
|
+
("neq", [True, True]),
|
|
198
|
+
],
|
|
199
|
+
)
|
|
200
|
+
def test_remaining_binary_contracts(self, name, expected):
|
|
201
|
+
df = pl.DataFrame({"a": [5, 7], "b": [2, 3]})
|
|
202
|
+
node = Call(name, (Column("a"), Column("b")))
|
|
203
|
+
|
|
204
|
+
assert df.select(ast_compile(node)).to_series().to_list() == expected
|
|
205
|
+
|
|
119
206
|
|
|
120
207
|
class TestHorizontalOps:
|
|
121
208
|
def test_max(self):
|
|
@@ -138,6 +225,17 @@ class TestHorizontalOps:
|
|
|
138
225
|
node = Call("mean", (Column("a"), Column("b")))
|
|
139
226
|
assert df.select(ast_compile(node))["mean(a, b)"].to_list() == [2.0, 4.0]
|
|
140
227
|
|
|
228
|
+
def test_arg_max_and_arg_min(self):
|
|
229
|
+
df = pl.DataFrame({"a": [1.0, 5.0], "b": [3.0, 2.0]})
|
|
230
|
+
|
|
231
|
+
result = df.select(
|
|
232
|
+
ast_compile(Call("arg_max", (Column("a"), Column("b")))).alias("max"),
|
|
233
|
+
ast_compile(Call("arg_min", (Column("a"), Column("b")))).alias("min"),
|
|
234
|
+
)
|
|
235
|
+
|
|
236
|
+
assert result["max"].to_list() == [1, 0]
|
|
237
|
+
assert result["min"].to_list() == [0, 1]
|
|
238
|
+
|
|
141
239
|
|
|
142
240
|
class TestTernary:
|
|
143
241
|
def test_if(self):
|
|
@@ -87,7 +87,9 @@ class TestCSQcut:
|
|
|
87
87
|
node = Call("cs_qcut", (Column("x"), Literal(2)))
|
|
88
88
|
expr = ast_compile(node)
|
|
89
89
|
result = cs_df.select(expr)
|
|
90
|
-
|
|
90
|
+
|
|
91
|
+
assert result.schema["cs_qcut(x, 2)"] == pl.Int32
|
|
92
|
+
assert result["cs_qcut(x, 2)"].to_list() == [1, 2, 1, 2, 2, 1]
|
|
91
93
|
|
|
92
94
|
|
|
93
95
|
class TestCSMidby:
|
|
@@ -96,3 +98,41 @@ class TestCSMidby:
|
|
|
96
98
|
expr = ast_compile(node)
|
|
97
99
|
result = cs_df.select(expr)
|
|
98
100
|
assert result.height == 6
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
class TestRemainingCSOperators:
|
|
104
|
+
@pytest.mark.parametrize(
|
|
105
|
+
("name", "expected_first"),
|
|
106
|
+
[
|
|
107
|
+
("cs_ufit", 5.0),
|
|
108
|
+
("cs_mid", 15.0),
|
|
109
|
+
("cs_corr", 1.0),
|
|
110
|
+
("cs_std", pytest.approx(7.0710678119)),
|
|
111
|
+
("cs_var", 50.0),
|
|
112
|
+
("cs_slope", 5.0),
|
|
113
|
+
("cs_resid", 5.0),
|
|
114
|
+
("cs_min", 10.0),
|
|
115
|
+
("cs_peakmax", False),
|
|
116
|
+
("cs_peakmin", False),
|
|
117
|
+
],
|
|
118
|
+
)
|
|
119
|
+
def test_operator_values(self, cs_df, name, expected_first):
|
|
120
|
+
args = (
|
|
121
|
+
(Column("x"), Column("y"))
|
|
122
|
+
if name in {"cs_corr", "cs_slope", "cs_resid"}
|
|
123
|
+
else (Column("x"),)
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
result = cs_df.select(ast_compile(Call(name, args)))
|
|
127
|
+
|
|
128
|
+
assert result.item(0, 0) == expected_first
|
|
129
|
+
|
|
130
|
+
def test_cs_skew_executes_per_cross_section(self, cs_df):
|
|
131
|
+
result = cs_df.select(ast_compile(Call("cs_skew", (Column("x"),))))
|
|
132
|
+
|
|
133
|
+
assert result.height == cs_df.height
|
|
134
|
+
|
|
135
|
+
def test_cs_meanby_adds_grouping_columns(self, cs_df):
|
|
136
|
+
result = cs_df.select(ast_compile(Call("cs_meanby", (Column("x"), Column("asset")))))
|
|
137
|
+
|
|
138
|
+
assert result.to_series().to_list() == cs_df["x"].to_list()
|
|
@@ -20,6 +20,12 @@ class TestCodeStrInit:
|
|
|
20
20
|
assert cs.data is None
|
|
21
21
|
assert "close" in cs.cache_columns
|
|
22
22
|
|
|
23
|
+
def test_init_collects_lazy_input_in_eager_mode(self, sample_df):
|
|
24
|
+
cs = CodeStr(sample_df.lazy())
|
|
25
|
+
|
|
26
|
+
assert isinstance(cs.data, pl.DataFrame)
|
|
27
|
+
assert cs.data.height == sample_df.height
|
|
28
|
+
|
|
23
29
|
def test_init_invalid_data_raises(self):
|
|
24
30
|
with pytest.raises(AssertionError):
|
|
25
31
|
CodeStr([1, 2, 3]) # type: ignore
|
|
@@ -83,6 +89,17 @@ class TestSQLInteractiveMode:
|
|
|
83
89
|
assert "total" in result.columns
|
|
84
90
|
assert "spread" in result.columns
|
|
85
91
|
|
|
92
|
+
def test_same_query_reuses_expression_with_new_alias(self, sample_df):
|
|
93
|
+
cs = CodeStr(sample_df)
|
|
94
|
+
|
|
95
|
+
result = cs.sql(
|
|
96
|
+
"close + 1 as first",
|
|
97
|
+
"close + 1 as second",
|
|
98
|
+
)
|
|
99
|
+
|
|
100
|
+
assert cs.failed == []
|
|
101
|
+
assert result["first"].equals(result["second"])
|
|
102
|
+
|
|
86
103
|
def test_sql_caching_reuses_expr(self, sample_df):
|
|
87
104
|
"""Subsequent identical expressions should reuse cached columns."""
|
|
88
105
|
cs = CodeStr(sample_df)
|
|
@@ -441,6 +458,116 @@ class TestCheckExpr:
|
|
|
441
458
|
assert result["reasons"] == []
|
|
442
459
|
|
|
443
460
|
|
|
461
|
+
class TestValidateExpr:
|
|
462
|
+
def test_validates_expression_against_current_schema(self, sample_df):
|
|
463
|
+
cs = CodeStr(sample_df)
|
|
464
|
+
|
|
465
|
+
assert cs.validate_expr("ts_mean(close, 2) as mean_close") == [
|
|
466
|
+
{
|
|
467
|
+
"expr": "ts_mean(close, 2) as mean_close",
|
|
468
|
+
"valid": True,
|
|
469
|
+
"stage": "schema",
|
|
470
|
+
"error_type": None,
|
|
471
|
+
"message": None,
|
|
472
|
+
}
|
|
473
|
+
]
|
|
474
|
+
|
|
475
|
+
def test_reports_structural_compile_and_schema_failures(self, sample_df):
|
|
476
|
+
cs = CodeStr(sample_df)
|
|
477
|
+
|
|
478
|
+
structural = cs.validate_expr("")[0]
|
|
479
|
+
compile_failure = cs.validate_expr("sin(1.0) as factor")[0]
|
|
480
|
+
schema_failure = cs.validate_expr("missing + 1 as factor")[0]
|
|
481
|
+
|
|
482
|
+
assert structural["valid"] is False
|
|
483
|
+
assert structural["stage"] == "structural"
|
|
484
|
+
assert structural["error_type"] == "ParseError"
|
|
485
|
+
assert structural["message"]
|
|
486
|
+
assert compile_failure["valid"] is False
|
|
487
|
+
assert compile_failure["stage"] == "compile"
|
|
488
|
+
assert compile_failure["error_type"] == "CompileError"
|
|
489
|
+
assert "sin" in compile_failure["message"]
|
|
490
|
+
assert schema_failure["valid"] is False
|
|
491
|
+
assert schema_failure["stage"] == "schema"
|
|
492
|
+
assert schema_failure["error_type"] == "ColumnNotFoundError"
|
|
493
|
+
assert "missing" in schema_failure["message"]
|
|
494
|
+
|
|
495
|
+
def test_reports_non_parse_structural_failure(self, sample_df):
|
|
496
|
+
cs = CodeStr(sample_df)
|
|
497
|
+
|
|
498
|
+
result = cs.validate_expr("close - close as invalid")[0]
|
|
499
|
+
|
|
500
|
+
assert result["valid"] is False
|
|
501
|
+
assert result["stage"] == "structural"
|
|
502
|
+
assert result["error_type"] == "StructuralError"
|
|
503
|
+
assert "redundant:sub" in result["message"]
|
|
504
|
+
|
|
505
|
+
def test_reports_invalid_base_lazy_schema(self):
|
|
506
|
+
invalid = pl.DataFrame({"x": [1.0]}).lazy().select("missing")
|
|
507
|
+
cs = CodeStr(invalid, pure_lazy=True)
|
|
508
|
+
|
|
509
|
+
result = cs.validate_expr("x + 1 as factor")[0]
|
|
510
|
+
|
|
511
|
+
assert result["valid"] is False
|
|
512
|
+
assert result["stage"] == "schema"
|
|
513
|
+
assert result["error_type"] == "ColumnNotFoundError"
|
|
514
|
+
|
|
515
|
+
def test_validates_mixed_window_through_planner(self, mixed_window_df):
|
|
516
|
+
cs = CodeStr(mixed_window_df, align=False)
|
|
517
|
+
|
|
518
|
+
result = cs.validate_expr("ts_mean(cs_moderate(x), 2, min_samples=1) as factor")
|
|
519
|
+
|
|
520
|
+
assert result[0]["valid"] is True
|
|
521
|
+
assert result[0]["stage"] == "schema"
|
|
522
|
+
|
|
523
|
+
def test_batch_expressions_are_independent(self, sample_df):
|
|
524
|
+
cs = CodeStr(sample_df)
|
|
525
|
+
|
|
526
|
+
results = cs.validate_expr(
|
|
527
|
+
"close * 2 as scaled",
|
|
528
|
+
"scaled + 1 as shifted",
|
|
529
|
+
)
|
|
530
|
+
|
|
531
|
+
assert results[0]["valid"] is True
|
|
532
|
+
assert results[1]["valid"] is False
|
|
533
|
+
assert results[1]["stage"] == "schema"
|
|
534
|
+
|
|
535
|
+
def test_does_not_collect_or_mutate_engine_state(self, sample_df):
|
|
536
|
+
cs = CodeStr(sample_df)
|
|
537
|
+
cs.sql("close + 1 as cached")
|
|
538
|
+
|
|
539
|
+
def raise_during_collect(value):
|
|
540
|
+
raise ValueError("must not collect")
|
|
541
|
+
|
|
542
|
+
def lazy_only(expr):
|
|
543
|
+
return expr.map_elements(
|
|
544
|
+
raise_during_collect,
|
|
545
|
+
return_dtype=pl.Float64,
|
|
546
|
+
)
|
|
547
|
+
|
|
548
|
+
cs.register_udf(lazy_only)
|
|
549
|
+
state = {
|
|
550
|
+
"data": cs.data,
|
|
551
|
+
"lazy": cs._data_,
|
|
552
|
+
"expr_cache": dict(cs._expr_cache),
|
|
553
|
+
"cur_expr_cache": dict(cs._cur_expr_cache),
|
|
554
|
+
"internal_columns": set(cs._internal_columns),
|
|
555
|
+
"last_query_cache": cs._last_query_cache,
|
|
556
|
+
"failed": list(cs.failed),
|
|
557
|
+
}
|
|
558
|
+
|
|
559
|
+
result = cs.validate_expr("lazy_only(close) as validated")
|
|
560
|
+
|
|
561
|
+
assert result[0]["valid"] is True
|
|
562
|
+
assert cs.data is state["data"]
|
|
563
|
+
assert cs._data_ is state["lazy"]
|
|
564
|
+
assert cs._expr_cache == state["expr_cache"]
|
|
565
|
+
assert cs._cur_expr_cache == state["cur_expr_cache"]
|
|
566
|
+
assert cs._internal_columns == state["internal_columns"]
|
|
567
|
+
assert cs._last_query_cache is state["last_query_cache"]
|
|
568
|
+
assert cs.failed == state["failed"]
|
|
569
|
+
|
|
570
|
+
|
|
444
571
|
class TestClearCache:
|
|
445
572
|
def test_clear_cache_resets_state(self, sample_df):
|
|
446
573
|
cs = CodeStr(sample_df)
|
|
@@ -51,7 +51,7 @@ wheels = [
|
|
|
51
51
|
|
|
52
52
|
[[package]]
|
|
53
53
|
name = "codestr"
|
|
54
|
-
version = "0.3.
|
|
54
|
+
version = "0.3.2"
|
|
55
55
|
source = { editable = "." }
|
|
56
56
|
dependencies = [
|
|
57
57
|
{ name = "lark" },
|
|
@@ -81,7 +81,7 @@ requires-dist = [
|
|
|
81
81
|
{ name = "loguru", specifier = ">=0.7,<1.0" },
|
|
82
82
|
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10" },
|
|
83
83
|
{ name = "numpy", specifier = ">=1.26,<3.0" },
|
|
84
|
-
{ name = "polars", specifier = ">=1.
|
|
84
|
+
{ name = "polars", specifier = ">=1.44.1" },
|
|
85
85
|
{ name = "pytest", marker = "extra == 'test'", specifier = ">=8.0" },
|
|
86
86
|
{ name = "pytest-cov", marker = "extra == 'test'", specifier = ">=5.0" },
|
|
87
87
|
{ name = "toolz", specifier = ">=1.0,<2.0" },
|
|
@@ -630,30 +630,30 @@ wheels = [
|
|
|
630
630
|
|
|
631
631
|
[[package]]
|
|
632
632
|
name = "polars"
|
|
633
|
-
version = "1.
|
|
633
|
+
version = "1.44.1"
|
|
634
634
|
source = { registry = "https://pypi.org/simple" }
|
|
635
635
|
dependencies = [
|
|
636
636
|
{ name = "polars-runtime-32" },
|
|
637
637
|
]
|
|
638
|
-
sdist = { url = "https://files.pythonhosted.org/packages/
|
|
638
|
+
sdist = { url = "https://files.pythonhosted.org/packages/26/73/258a1fe17bb2744a507199566ed712663144fdd0811b615b59a47dfa38d2/polars-1.44.1.tar.gz", hash = "sha256:ef3c89e9ebbbe8eb343c06873f1945683f8b6f97a1bdf001c60551c6c5e3cda1", size = 765660, upload-time = "2026-08-26T07:09:12.704Z" }
|
|
639
639
|
wheels = [
|
|
640
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
640
|
+
{ url = "https://files.pythonhosted.org/packages/3a/f1/59154659081930fbde291ea6225607956b500a71b0ce45d88217b7d32da2/polars-1.44.1-py3-none-any.whl", hash = "sha256:1fa62fc1c88fba77a68b28291b5aabdd69e5f38b34e59721a064ae3169b59bb5", size = 865208, upload-time = "2026-08-26T07:07:44.646Z" },
|
|
641
641
|
]
|
|
642
642
|
|
|
643
643
|
[[package]]
|
|
644
644
|
name = "polars-runtime-32"
|
|
645
|
-
version = "1.
|
|
645
|
+
version = "1.44.1"
|
|
646
646
|
source = { registry = "https://pypi.org/simple" }
|
|
647
|
-
sdist = { url = "https://files.pythonhosted.org/packages/
|
|
647
|
+
sdist = { url = "https://files.pythonhosted.org/packages/fd/b2/2a76415d047a45df05489f2334c91ff120a274cf655d4ca030c7f54a8743/polars_runtime_32-1.44.1.tar.gz", hash = "sha256:abd10a54ed1caff42228610fcba0f93251f9870bd7cffb0c78bc26f5e0718ce4", size = 3171156, upload-time = "2026-08-26T07:09:14.243Z" }
|
|
648
648
|
wheels = [
|
|
649
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
650
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
651
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
652
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
653
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
654
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
655
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
656
|
-
{ url = "https://files.pythonhosted.org/packages/
|
|
649
|
+
{ url = "https://files.pythonhosted.org/packages/f1/93/ef9344dcec16757cf21027dc907ef989197e50742cfe4407f2e87edb0a7f/polars_runtime_32-1.44.1-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:1dfccb2b52aa50468a7d28e3e61c8338a13fb5bffc8646e388a649f5bdc6b463", size = 53962370, upload-time = "2026-08-26T07:07:47.21Z" },
|
|
650
|
+
{ url = "https://files.pythonhosted.org/packages/10/da/38b32b7901af33f1fee2172ceaa39e9159825657920064298917392d78fa/polars_runtime_32-1.44.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:0580807dc3eed258f0db70bb65d905dd43f0135392119ec25308033ae24258fb", size = 48620468, upload-time = "2026-08-26T07:07:50.436Z" },
|
|
651
|
+
{ url = "https://files.pythonhosted.org/packages/49/51/185af877d1d2236671493cf72bd3327a6046240eeea69c0696d1af2a5acb/polars_runtime_32-1.44.1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0627f9aa82cb869725235e5188f698862fd9ada0c8c1cf65c3dc5a49a4a0ec26", size = 52455947, upload-time = "2026-08-26T07:07:53.776Z" },
|
|
652
|
+
{ url = "https://files.pythonhosted.org/packages/fc/0a/0858f60cb5a6f8f73ec4cdd73eccd9f748d66bfc23304c5c23fa3468094a/polars_runtime_32-1.44.1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eea4283be8e60822d890dbda20588fe59b4172b508bd5ebf3471e531ca9f50d7", size = 58561262, upload-time = "2026-08-26T07:07:57.508Z" },
|
|
653
|
+
{ url = "https://files.pythonhosted.org/packages/73/08/de4774b5612d7c8739f89ac01b601486b4f057b1da35a5b876bf9276fd95/polars_runtime_32-1.44.1-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:04e2c0f46e7a9906fffb1897f18f23b079b74f83c56b50060bace9e7b9b49b1a", size = 52636064, upload-time = "2026-08-26T07:08:06.337Z" },
|
|
654
|
+
{ url = "https://files.pythonhosted.org/packages/98/ac/769c598dd106e2a6647798da3ed25ddeee67f2d12c04f5da316cb3da6360/polars_runtime_32-1.44.1-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:0956f0cae632d8fad3a04b4315bf2bb69b56d10c83c79a75c2c4c5a13b9ce5cc", size = 56447612, upload-time = "2026-08-26T07:08:12.36Z" },
|
|
655
|
+
{ url = "https://files.pythonhosted.org/packages/b2/ee/98408296e15388020b6183323fdbe78ccab4f72c20d8e0d7092c062d3ad2/polars_runtime_32-1.44.1-cp310-abi3-win_amd64.whl", hash = "sha256:159334184e6fbb074c9f4692221ea19970a5e2bed2a479f9d7bdb00b7f3eedb9", size = 53702970, upload-time = "2026-08-26T07:08:15.398Z" },
|
|
656
|
+
{ url = "https://files.pythonhosted.org/packages/e9/9d/8b17e075aac73c881a50b6c1f690d20df46db2f3bcabc99f600ecdee1290/polars_runtime_32-1.44.1-cp310-abi3-win_arm64.whl", hash = "sha256:3ba28d638d0513e0b4afbcdab5c0059a85021e5f81d62b5f793e7e23badb2cf7", size = 47281050, upload-time = "2026-08-26T07:08:18.43Z" },
|
|
657
657
|
]
|
|
658
658
|
|
|
659
659
|
[[package]]
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/plans/2026-07-16-ts-ema-keyword-arguments.md
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
{codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-16-polars-minimum-version-design.md
RENAMED
|
File without changes
|
{codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-16-ts-ema-keyword-arguments-design.md
RENAMED
|
File without changes
|
|
File without changes
|
{codestr-0.3.0 → codestr-0.3.2}/docs/superpowers/specs/2026-07-17-mixed-window-planner-design.md
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|