codestr 0.1.0__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.
@@ -0,0 +1,35 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [master]
6
+ pull_request:
7
+ branches: [master]
8
+
9
+ jobs:
10
+ lint:
11
+ runs-on: ubuntu-latest
12
+ steps:
13
+ - uses: actions/checkout@v4
14
+ - uses: astral-sh/setup-uv@v5
15
+ with:
16
+ python-version: "3.12"
17
+ - name: ruff check
18
+ run: uvx ruff check src/
19
+ - name: ruff format check
20
+ run: uvx ruff format --check src/
21
+
22
+ test:
23
+ runs-on: ubuntu-latest
24
+ strategy:
25
+ matrix:
26
+ python-version: ["3.10", "3.11", "3.12"]
27
+ steps:
28
+ - uses: actions/checkout@v4
29
+ - uses: astral-sh/setup-uv@v5
30
+ with:
31
+ python-version: ${{ matrix.python-version }}
32
+ - name: Install and test
33
+ run: |
34
+ uv sync --extra test
35
+ uv run pytest tests/ -v
@@ -0,0 +1,41 @@
1
+ name: Release
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+
7
+ jobs:
8
+ build:
9
+ runs-on: ubuntu-latest
10
+ permissions:
11
+ contents: read
12
+ steps:
13
+ - uses: actions/checkout@v4
14
+ - uses: astral-sh/setup-uv@v5
15
+ with:
16
+ python-version: "3.12"
17
+ - name: Build distributions
18
+ run: uv build
19
+ - name: Upload distributions
20
+ uses: actions/upload-artifact@v4
21
+ with:
22
+ name: dist
23
+ path: dist/
24
+
25
+ publish:
26
+ needs: build
27
+ runs-on: ubuntu-latest
28
+ environment:
29
+ name: pypi
30
+ url: https://pypi.org/p/codestr
31
+ permissions:
32
+ contents: read
33
+ id-token: write
34
+ steps:
35
+ - name: Download distributions
36
+ uses: actions/download-artifact@v4
37
+ with:
38
+ name: dist
39
+ path: dist
40
+ - name: Publish package distributions to PyPI
41
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,7 @@
1
+ __pycache__/
2
+ .venv/
3
+ *.pyc
4
+ .idea/
5
+ .vscode/
6
+ *.egg-info/
7
+ dist/
@@ -0,0 +1 @@
1
+ 3.12
@@ -0,0 +1,37 @@
1
+ # Changelog
2
+
3
+ ## [Unreleased]
4
+
5
+ ### Added
6
+ - **Per-instance over-window config**: `partition_by` / `order_by` params on `CodeStr.__init__`
7
+ replacing the global `over` dict and `configure_over()`.
8
+ - Multi-column support for `partition_by` and `order_by`.
9
+ - Compiler auto-injects `partition_by`/`order_by` based on `UDFMeta.category`.
10
+ - 163 tests covering parser, compiler, engine, syntax, registry, and all UDF operators.
11
+
12
+ ### Changed
13
+ - `CodeStr.__init__` signature: `index` is now a 2-tuple `(time_col, entity_col)`.
14
+ - TS/CS UDF functions accept `partition_by`/`order_by` keyword arguments (injected by compiler).
15
+ - `udf/__init__.py` uses explicit module imports instead of `import *`.
16
+ - Stricter ruff rules: added B, SIM, RUF, TCH.
17
+
18
+ ### Removed
19
+ - **BREAKING**: Module-level `over` dict and `configure_over()` from `ts_udf` and `cs_udf`.
20
+
21
+ ### Fixed
22
+ - `cs_midby`/`cs_meanby`: `[*over, *by]` was unpacking dict keys as column names —
23
+ now uses `partition_by + list(by)`.
24
+
25
+ ## [0.1.0] — Initial Release
26
+
27
+ ### Added
28
+ - DSL parser (Lark LALR grammar): binary/unary/ternary operators, function calls, implicit multiplication
29
+ - AST compiler that translates expressions to Polars expressions
30
+ - `CodeStr` engine with two API modes: pure `compile()` and interactive `sql()`
31
+ - Expression-level caching with hash-based deduplication (alias-independent)
32
+ - Built-in UDF operators:
33
+ - **base_udf**: arithmetic, logical, trigonometric, and horizontal operations
34
+ - **cs_udf**: cross-section operators (rank, z-score, IC, quantile cut, etc.)
35
+ - **ts_udf**: time-series operators (rolling mean, sum, delay, delta, etc.)
36
+ - UDF registry with `@udf` decorator for custom operators
37
+ - CI workflow with lint and format checks (ruff)
@@ -0,0 +1,105 @@
1
+ # CLAUDE.md
2
+
3
+ This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
4
+
5
+ ## 项目概述
6
+
7
+ CodeStr 是一个专为量化因子挖掘设计的表达式计算引擎,基于 [Polars](https://pola.rs/) 构建,提供 DSL → Polars Expr 的高效转译与执行。
8
+
9
+ ## 开发命令
10
+
11
+ ```bash
12
+ uv sync --extra dev # 安装项目 + 开发依赖
13
+ uv run pytest tests/ -v # 运行所有测试
14
+ uv run pytest tests/test_parser.py -v # 运行单个测试文件
15
+
16
+ # 代码质量(与 CI 一致)
17
+ uvx ruff check src/ tests/ # lint 检查
18
+ uvx ruff format --check src/ # 格式检查(CI 模式)
19
+ uvx ruff format src/ # 自动格式化
20
+ ```
21
+
22
+ 项目使用 `hatchling` 构建,`ruff` 做 lint/格式化,`mypy` 做类型检查。CI 见 `.github/workflows/ci.yml`,在 push/PR 到 master 时运行 lint + format + test matrix (3.10/3.11/3.12)。
23
+
24
+ ## 项目结构
25
+
26
+ ```
27
+ src/codestr/
28
+ ├── __init__.py # 公开 API:CodeStr, ExprNode, Call, Column, Literal
29
+ ├── engine.py # CodeStr 引擎:compile() 纯编译 + sql() 交互式有状态
30
+ ├── compiler.py # AST → Polars Expr 编译器(纯函数,无副作用)
31
+ ├── parser.py # DSL 解析器(Lark LALR grammar + Transformer → AST)
32
+ ├── syntax.py # AST 定义:ExprNode, Column, Literal, Call + 分析辅助函数
33
+ ├── tokens.py # Token/TokenType 定义(RPN 分析用)
34
+ ├── errors.py # 异常:ParseError, CompileError, PolarsError, FailError
35
+ └── udf/
36
+ ├── registry.py # UDF 注册中心(单例 + @udf 装饰器)
37
+ ├── base_udf.py # 基础算子:一元/二元/三元 + 算术/逻辑运算
38
+ ├── cs_udf.py # 截面算子 (Cross-Section),如 cs_rank, cs_zscore
39
+ └── ts_udf.py # 时序算子 (Time-Series),如 ts_mean, ts_delay
40
+ ```
41
+
42
+ ## 核心架构
43
+
44
+ ### 数据流
45
+
46
+ ```
47
+ DSL 字符串 → parser.parse() → AST (ExprNode) → compiler.compile() → pl.Expr → Polars 执行
48
+ ```
49
+
50
+ - **Parser** (`parser.py`): 用 Lark LALR 解析器将 DSL 字符串转为 AST。`_normalize()` 预处理:`if(` → `if_(`,`!` → `~`(但保留 `!=`),移除 `$` 和换行。`as` 关键字后的别名会被注入 AST 的 `_alias` 字段。
51
+ - **Compiler** (`compiler.py`): 纯函数,AST → `pl.Expr`。`_resolve()` 对 Literal 返回 Python 标量,对 Column/Call 返回 `pl.Expr`。`_compile()` 总是返回 `pl.Expr`(Literal 包装为 `pl.lit()`)。函数从 UDFRegistry 查找。引擎上下文通过签名检查自动注入:`dims`、`partition_by`、`order_by`。
52
+ - **Engine** (`engine.py`): `CodeStr` 类封装了上述流程,额外管理状态(数据对齐、表达式缓存、惰性计算图)。
53
+
54
+ ### 两种 API 模式
55
+
56
+ | 模式 | API | 副作用 | 适用场景 |
57
+ |------|-----|--------|---------|
58
+ | 纯编译 | `compile(expr_str) -> pl.Expr` | 无 | 批量/RL 评估 |
59
+ | 交互式 | `sql(expr_str, lazy=False) -> pl.DataFrame` | 有(更新内部状态) | 交互式研究/GP |
60
+
61
+ - `sql(lazy=True)` 返回 `pl.LazyFrame`,不物化,不更新内部状态(rollback `_data_`)。
62
+ - `sql(lazy=False)`(默认)collect 后物化,更新 `_expr_cache` 和 `data`。
63
+
64
+ ### 窗口配置
65
+
66
+ CodeStr 通过 `partition_by` / `order_by` 参数配置算子窗口,**不再有全局 `over` 字典**:
67
+
68
+ ```python
69
+ CodeStr(df,
70
+ index=("datetime", "asset"), # 面板对齐 + 结果选择
71
+ partition_by=["asset", "industry"], # 实体分组轴
72
+ order_by=["datetime", "tick"], # 时间排序轴
73
+ )
74
+ ```
75
+
76
+ - **TS 算子**: `over(partition_by=partition_by, order_by=order_by)` — 按实体分组,沿时间排序
77
+ - **CS 算子**: `over(partition_by=order_by, order_by=partition_by)` — 按时间分组,沿实体排序(交换)
78
+
79
+ 编译器在编译期根据 `UDFMeta.category` 自动注入对应的窗口配置。`category="math"` / `"user"` 不注入。
80
+
81
+ ### 表达式缓存策略
82
+
83
+ 两层缓存,都基于 `Call.__hash__`(只看 `fn_name + args`,**别名不参与**):
84
+
85
+ 1. **`_expr_cache`** — 跨查询持久化缓存,只在 eager `sql()` 成功后合并入。
86
+ 2. **`_cur_expr_cache`** — 单次查询临时缓存,查询成功后才 merge 到 `_expr_cache`。
87
+
88
+ 缓存命中时直接 `pl.col(cached_alias).alias(new_alias)` 复用,避免重编译。`cover=True` 强制绕过缓存。
89
+
90
+ ### UDF 注册
91
+
92
+ `UDFRegistry` 是单例。三种注册方式:
93
+ 1. `@udf(category="ts")` 装饰器(自动推断 arity)
94
+ 2. `CodeStr.register_udf(func, name="...")` 实例方法
95
+ 3. 直接 `UDFRegistry.get_instance().register(UDFMeta(...))`
96
+
97
+ 编译器通过 `inspect.signature` 检查函数签名,自动注入:
98
+ - `dims` — 如果函数接受且引擎有 dims 信息
99
+ - `partition_by` / `order_by` — 根据 category(ts → ts_over, cs → cs_over)
100
+
101
+ ### DSL 语法
102
+
103
+ 运算符优先级(从低到高):三元 `?:` → 或 `|` → 与 `&` → 比较 `<>` `<=` `>=` `==` `!=` → 加减 `+` `-` → 乘除 `*` `/` `//` `%` → 幂 `**` → 一元 `-` `~` → 函数调用/属性访问
104
+
105
+ 支持隐式乘法:`5close` → `5 * close`。支持属性链:`df.close` → `Column("df.close")`。
@@ -0,0 +1,62 @@
1
+ # Contributing to CodeStr
2
+
3
+ CodeStr is an expression compute engine for quantitative factor mining — contributions are welcome.
4
+
5
+ ## Getting started
6
+
7
+ ```bash
8
+ git clone https://github.com/huangbogeng/codestr.git
9
+ cd codestr
10
+ uv sync --extra dev
11
+ ```
12
+
13
+ ## Development workflow
14
+
15
+ ```bash
16
+ # 1. Create a feature branch from master
17
+ git checkout -b feature/my-change
18
+
19
+ # 2. Make your changes
20
+
21
+ # 3. Run checks locally (same as CI)
22
+ uvx ruff check src/ tests/
23
+ uvx ruff format --check src/ tests/
24
+ uv run pytest tests/ -v
25
+
26
+ # 4. Commit and push
27
+ git commit -m "feat: describe your change"
28
+ git push -u origin feature/my-change
29
+ ```
30
+
31
+ ## Code style
32
+
33
+ - **Formatter**: [Ruff](https://docs.astral.sh/ruff/) — 100 char line length, double quotes
34
+ - **Type checker**: [mypy](https://www.mypy-lang.org/) — Python 3.10 target
35
+ - **Linter**: Ruff with E, F, I, N, W, UP, B, SIM, RUF, TCH rules
36
+ - **Tests**: pytest in `tests/` directory
37
+ - **Build**: [hatchling](https://hatch.pypa.io/)
38
+
39
+ CI will fail if any check doesn't pass.
40
+
41
+ ## Testing
42
+
43
+ - Test files mirror the source structure: `tests/test_engine.py` tests `src/codestr/engine.py`
44
+ - Use `conftest.py` fixtures for shared test data
45
+ - The `reset_registry` fixture (autouse) resets the UDF registry between tests — register new UDFs in test functions, not at module level
46
+
47
+ ```bash
48
+ uv run pytest tests/ # all tests
49
+ uv run pytest tests/test_parser.py -v # single file
50
+ uv run pytest tests/ -k "test_ts" -v # filter by name
51
+ ```
52
+
53
+ ## Project architecture
54
+
55
+ See [CLAUDE.md](./CLAUDE.md) for a detailed walkthrough of the data flow, caching, and window configuration.
56
+
57
+ ## Reporting issues
58
+
59
+ Please include:
60
+ - A minimal reproducible example
61
+ - Expected vs actual behavior
62
+ - Python version (`python --version`) and package versions (`uv pip list | grep -E "polars|codestr|lark"`)
codestr-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2025 huangbogeng
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
codestr-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,157 @@
1
+ Metadata-Version: 2.4
2
+ Name: codestr
3
+ Version: 0.1.0
4
+ Summary: DSL → Polars Expr compile engine for quantitative factor mining
5
+ Project-URL: Repository, https://github.com/huangbogeng/codestr
6
+ Author: huangbogeng
7
+ License: MIT
8
+ License-File: LICENSE
9
+ Keywords: dsl,expression-engine,factor-mining,polars,quantitative-finance
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Financial and Insurance Industry
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: 3.12
17
+ Requires-Python: >=3.10
18
+ Requires-Dist: lark<2.0,>=1.0
19
+ Requires-Dist: loguru<1.0,>=0.7
20
+ Requires-Dist: numpy<3.0,>=1.26
21
+ Requires-Dist: polars<2.0,>=1.0
22
+ Requires-Dist: toolz<2.0,>=1.0
23
+ Provides-Extra: dev
24
+ Requires-Dist: mypy>=1.10; extra == 'dev'
25
+ Requires-Dist: pytest-cov>=5.0; extra == 'dev'
26
+ Requires-Dist: pytest>=8.0; extra == 'dev'
27
+ Provides-Extra: test
28
+ Requires-Dist: pytest-cov>=5.0; extra == 'test'
29
+ Requires-Dist: pytest>=8.0; extra == 'test'
30
+ Description-Content-Type: text/markdown
31
+
32
+ # CodeStr
33
+
34
+ [![CI](https://github.com/huangbogeng/codestr/actions/workflows/ci.yml/badge.svg)](https://github.com/huangbogeng/codestr/actions/workflows/ci.yml)
35
+ [![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](https://www.python.org/)
36
+ [![License](https://img.shields.io/badge/license-MIT-green)](./LICENSE)
37
+ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
38
+
39
+ CodeStr 是一个专为量化因子挖掘设计的 DSL → Polars Expr 表达式计算引擎,提供高效的表达式转译、缓存与执行。
40
+
41
+ ## 安装
42
+
43
+ ```bash
44
+ git clone https://github.com/huangbogeng/codestr.git
45
+ cd codestr
46
+ uv sync --extra dev
47
+ ```
48
+
49
+ ## 快速开始
50
+
51
+ ```python
52
+ import polars as pl
53
+ from codestr import CodeStr
54
+
55
+ # 标准面板数据 (time, entity)
56
+ df = pl.DataFrame({
57
+ "datetime": ["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"],
58
+ "asset": ["A", "B", "A", "B"],
59
+ "close": [100.0, 200.0, 101.0, 198.0],
60
+ "volume": [1000.0, 2000.0, 1100.0, 1900.0],
61
+ })
62
+
63
+ cs = CodeStr(df, index=("datetime", "asset"))
64
+
65
+ # 交互式查询 — 结果自动缓存
66
+ result = cs.sql(
67
+ "ts_mean(close, 5) as ma5",
68
+ "cs_rank(close) as rank",
69
+ "close / ts_delay(close, 1) - 1 as ret",
70
+ )
71
+ print(result)
72
+ ```
73
+
74
+ ## 两种 API 模式
75
+
76
+ | 模式 | API | 行为 |
77
+ |------|-----|------|
78
+ | **纯编译** | `cs.compile(expr) -> pl.Expr` | 无副作用,返回 Polars 表达式 |
79
+ | **交互式** | `cs.sql(expr, lazy=False) -> pl.DataFrame` | 有状态,自动缓存与复用 |
80
+
81
+ ```python
82
+ # 纯编译 — 表达式可被任意 DataFrame 消费
83
+ expr = cs.compile("ts_mean(close, 5) as ma5")
84
+ other_df.with_columns(expr)
85
+
86
+ # 交互式 — 适合逐步构建因子
87
+ cs.sql("close + volume as total")
88
+ cs.sql("ts_mean(total, 5) as total_ma5") # 复用上一步的 total
89
+ ```
90
+
91
+ ## 窗口配置
92
+
93
+ CodeStr 使用 `partition_by`(实体分组轴)和 `order_by`(时间排序轴)控制窗口算子:
94
+
95
+ ```python
96
+ # 默认配置
97
+ cs = CodeStr(df)
98
+ # index=("datetime", "asset")
99
+ # → TS: over(partition_by=["asset"], order_by=["datetime"])
100
+ # → CS: over(partition_by=["datetime"], order_by=["asset"])
101
+
102
+ # 自定义列名
103
+ cs = CodeStr(df, index=("trade_date", "stock_code"))
104
+
105
+ # 多列窗口 — 按行业+股票分组,按日期+逐笔序号排序
106
+ cs = CodeStr(df,
107
+ index=("trade_date", "stock_code"),
108
+ partition_by=["industry", "stock_code"],
109
+ order_by=["trade_date", "tick"],
110
+ )
111
+ ```
112
+
113
+ | 算子类别 | 窗口规则 |
114
+ |---------|---------|
115
+ | **TS (时序)** | `over(partition_by=partition_by, order_by=order_by)` |
116
+ | **CS (截面)** | `over(partition_by=order_by, order_by=partition_by)` |
117
+
118
+ ## 自定义算子
119
+
120
+ ```python
121
+ from codestr.udf.registry import udf
122
+ import polars as pl
123
+
124
+ @udf(category="ts")
125
+ def ts_ewm(expr: pl.Expr, windows, partition_by=None, order_by=None):
126
+ """指数加权移动平均"""
127
+ return expr.ewm_mean(halflife=windows).over(
128
+ partition_by=partition_by, order_by=order_by
129
+ )
130
+
131
+ cs.sql("ts_ewm(close, 10) as ewm10")
132
+ ```
133
+
134
+ ## 内置算子
135
+
136
+ **基础算子** (`base_udf`):`abs`, `log`, `sqrt`, `square`, `cube`, `sin`, `cos`, `tan`, `exp`, `sigmoid`, `sign`, `clip`, `trunc`, `between`, `cast`, `max`, `min`, `sum`, `mean`, `arg_max`, `arg_min`, `if_`, `fib` 等
137
+
138
+ **截面算子** (`cs_udf`):`cs_rank`, `cs_zscore`, `cs_demean`, `cs_mean`, `cs_std`, `cs_var`, `cs_skew`, `cs_ic`, `cs_corr`, `cs_slope`, `cs_resid`, `cs_qcut`, `cs_midby`, `cs_meanby` 等
139
+
140
+ **时序算子** (`ts_udf`):`ts_mean`, `ts_sum`, `ts_std`, `ts_var`, `ts_skew`, `ts_kurt`, `ts_max`, `ts_min`, `ts_mid`, `ts_delay`, `ts_delta`, `ts_mad` 等
141
+
142
+ ## 项目结构
143
+
144
+ ```
145
+ src/codestr/
146
+ ├── engine.py # CodeStr 引擎入口
147
+ ├── compiler.py # AST → Polars Expr 编译器
148
+ ├── parser.py # DSL 解析器 (Lark LALR grammar)
149
+ ├── syntax.py # AST 节点定义
150
+ ├── tokens.py # Token 定义
151
+ ├── errors.py # 异常类型
152
+ └── udf/
153
+ ├── registry.py # UDF 注册中心 (@udf 装饰器)
154
+ ├── base_udf.py # 基础算子
155
+ ├── cs_udf.py # 截面算子 (Cross-Section)
156
+ └── ts_udf.py # 时序算子 (Time-Series)
157
+ ```
@@ -0,0 +1,126 @@
1
+ # CodeStr
2
+
3
+ [![CI](https://github.com/huangbogeng/codestr/actions/workflows/ci.yml/badge.svg)](https://github.com/huangbogeng/codestr/actions/workflows/ci.yml)
4
+ [![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](https://www.python.org/)
5
+ [![License](https://img.shields.io/badge/license-MIT-green)](./LICENSE)
6
+ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
7
+
8
+ CodeStr 是一个专为量化因子挖掘设计的 DSL → Polars Expr 表达式计算引擎,提供高效的表达式转译、缓存与执行。
9
+
10
+ ## 安装
11
+
12
+ ```bash
13
+ git clone https://github.com/huangbogeng/codestr.git
14
+ cd codestr
15
+ uv sync --extra dev
16
+ ```
17
+
18
+ ## 快速开始
19
+
20
+ ```python
21
+ import polars as pl
22
+ from codestr import CodeStr
23
+
24
+ # 标准面板数据 (time, entity)
25
+ df = pl.DataFrame({
26
+ "datetime": ["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"],
27
+ "asset": ["A", "B", "A", "B"],
28
+ "close": [100.0, 200.0, 101.0, 198.0],
29
+ "volume": [1000.0, 2000.0, 1100.0, 1900.0],
30
+ })
31
+
32
+ cs = CodeStr(df, index=("datetime", "asset"))
33
+
34
+ # 交互式查询 — 结果自动缓存
35
+ result = cs.sql(
36
+ "ts_mean(close, 5) as ma5",
37
+ "cs_rank(close) as rank",
38
+ "close / ts_delay(close, 1) - 1 as ret",
39
+ )
40
+ print(result)
41
+ ```
42
+
43
+ ## 两种 API 模式
44
+
45
+ | 模式 | API | 行为 |
46
+ |------|-----|------|
47
+ | **纯编译** | `cs.compile(expr) -> pl.Expr` | 无副作用,返回 Polars 表达式 |
48
+ | **交互式** | `cs.sql(expr, lazy=False) -> pl.DataFrame` | 有状态,自动缓存与复用 |
49
+
50
+ ```python
51
+ # 纯编译 — 表达式可被任意 DataFrame 消费
52
+ expr = cs.compile("ts_mean(close, 5) as ma5")
53
+ other_df.with_columns(expr)
54
+
55
+ # 交互式 — 适合逐步构建因子
56
+ cs.sql("close + volume as total")
57
+ cs.sql("ts_mean(total, 5) as total_ma5") # 复用上一步的 total
58
+ ```
59
+
60
+ ## 窗口配置
61
+
62
+ CodeStr 使用 `partition_by`(实体分组轴)和 `order_by`(时间排序轴)控制窗口算子:
63
+
64
+ ```python
65
+ # 默认配置
66
+ cs = CodeStr(df)
67
+ # index=("datetime", "asset")
68
+ # → TS: over(partition_by=["asset"], order_by=["datetime"])
69
+ # → CS: over(partition_by=["datetime"], order_by=["asset"])
70
+
71
+ # 自定义列名
72
+ cs = CodeStr(df, index=("trade_date", "stock_code"))
73
+
74
+ # 多列窗口 — 按行业+股票分组,按日期+逐笔序号排序
75
+ cs = CodeStr(df,
76
+ index=("trade_date", "stock_code"),
77
+ partition_by=["industry", "stock_code"],
78
+ order_by=["trade_date", "tick"],
79
+ )
80
+ ```
81
+
82
+ | 算子类别 | 窗口规则 |
83
+ |---------|---------|
84
+ | **TS (时序)** | `over(partition_by=partition_by, order_by=order_by)` |
85
+ | **CS (截面)** | `over(partition_by=order_by, order_by=partition_by)` |
86
+
87
+ ## 自定义算子
88
+
89
+ ```python
90
+ from codestr.udf.registry import udf
91
+ import polars as pl
92
+
93
+ @udf(category="ts")
94
+ def ts_ewm(expr: pl.Expr, windows, partition_by=None, order_by=None):
95
+ """指数加权移动平均"""
96
+ return expr.ewm_mean(halflife=windows).over(
97
+ partition_by=partition_by, order_by=order_by
98
+ )
99
+
100
+ cs.sql("ts_ewm(close, 10) as ewm10")
101
+ ```
102
+
103
+ ## 内置算子
104
+
105
+ **基础算子** (`base_udf`):`abs`, `log`, `sqrt`, `square`, `cube`, `sin`, `cos`, `tan`, `exp`, `sigmoid`, `sign`, `clip`, `trunc`, `between`, `cast`, `max`, `min`, `sum`, `mean`, `arg_max`, `arg_min`, `if_`, `fib` 等
106
+
107
+ **截面算子** (`cs_udf`):`cs_rank`, `cs_zscore`, `cs_demean`, `cs_mean`, `cs_std`, `cs_var`, `cs_skew`, `cs_ic`, `cs_corr`, `cs_slope`, `cs_resid`, `cs_qcut`, `cs_midby`, `cs_meanby` 等
108
+
109
+ **时序算子** (`ts_udf`):`ts_mean`, `ts_sum`, `ts_std`, `ts_var`, `ts_skew`, `ts_kurt`, `ts_max`, `ts_min`, `ts_mid`, `ts_delay`, `ts_delta`, `ts_mad` 等
110
+
111
+ ## 项目结构
112
+
113
+ ```
114
+ src/codestr/
115
+ ├── engine.py # CodeStr 引擎入口
116
+ ├── compiler.py # AST → Polars Expr 编译器
117
+ ├── parser.py # DSL 解析器 (Lark LALR grammar)
118
+ ├── syntax.py # AST 节点定义
119
+ ├── tokens.py # Token 定义
120
+ ├── errors.py # 异常类型
121
+ └── udf/
122
+ ├── registry.py # UDF 注册中心 (@udf 装饰器)
123
+ ├── base_udf.py # 基础算子
124
+ ├── cs_udf.py # 截面算子 (Cross-Section)
125
+ └── ts_udf.py # 时序算子 (Time-Series)
126
+ ```