probnode-sdk 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.
- probnode_sdk-0.1.0/PKG-INFO +50 -0
- probnode_sdk-0.1.0/README.md +30 -0
- probnode_sdk-0.1.0/pyproject.toml +43 -0
- probnode_sdk-0.1.0/src/probnode/__init__.py +7 -0
- probnode_sdk-0.1.0/src/probnode/cli.py +227 -0
- probnode_sdk-0.1.0/src/probnode/client.py +289 -0
- probnode_sdk-0.1.0/src/probnode/lp2.py +357 -0
- probnode_sdk-0.1.0/src/probnode/read.py +47 -0
- probnode_sdk-0.1.0/tests/test_client_cli.py +204 -0
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: probnode-sdk
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Python client and CLI for ProbNode: Polymarket order books, trades and gaps
|
|
5
|
+
Project-URL: Homepage, https://probnode.com
|
|
6
|
+
Project-URL: Documentation, https://probnode.com/docs
|
|
7
|
+
License: MIT
|
|
8
|
+
Requires-Python: >=3.10
|
|
9
|
+
Requires-Dist: httpx>=0.27
|
|
10
|
+
Requires-Dist: numpy>=1.24
|
|
11
|
+
Requires-Dist: pyarrow>=14
|
|
12
|
+
Provides-Extra: dev
|
|
13
|
+
Requires-Dist: pytest>=8.3; extra == 'dev'
|
|
14
|
+
Requires-Dist: ruff>=0.6; extra == 'dev'
|
|
15
|
+
Provides-Extra: pandas
|
|
16
|
+
Requires-Dist: pandas>=2.0; extra == 'pandas'
|
|
17
|
+
Provides-Extra: polars
|
|
18
|
+
Requires-Dist: polars>=0.20; extra == 'polars'
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
|
|
21
|
+
# probnode-sdk
|
|
22
|
+
|
|
23
|
+
Python client and command line for [ProbNode](https://probnode.com): Polymarket order books, trades
|
|
24
|
+
and feed gaps as hourly and market-day files (lp2, parquet, csv.gz), with a sha256 for every file.
|
|
25
|
+
|
|
26
|
+
```sh
|
|
27
|
+
pip install ./packages/probnode # from a checkout; on PyPI as probnode-sdk soon
|
|
28
|
+
# (not `pip install probnode`: that is someone else's package)
|
|
29
|
+
probnode login
|
|
30
|
+
probnode datasets
|
|
31
|
+
probnode files books parquet 2026-11-08T12 2026-11-08T14
|
|
32
|
+
probnode sync books parquet 2026-11-08T12 2026-11-08T14 -o data/ # again later: only changes
|
|
33
|
+
probnode doctor data/ # re-check every sha256
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
import probnode
|
|
38
|
+
|
|
39
|
+
pn = probnode.Client() # PROBNODE_API_KEY, or the key `probnode login` saved
|
|
40
|
+
files = pn.files("books", "parquet", "2026-11-08T12:00:00Z", "2026-11-08T13:00:00Z")
|
|
41
|
+
paths = list(pn.download_all(files, "data/")) # each checked against its sha256
|
|
42
|
+
table = probnode.read(paths[0]) # pyarrow.Table -> .to_pandas() / polars.from_arrow()
|
|
43
|
+
|
|
44
|
+
for batch in probnode.iter_lp2("hour.lp2"): # native books, decoded in bounded memory
|
|
45
|
+
...
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
Units: prices x10,000 (5400 = 0.54), sizes x1,000,000 shares, timestamps in ms UTC. Signed links
|
|
49
|
+
last 15 minutes; the client signs again whenever it needs to. The import name is `probnode`; the
|
|
50
|
+
PyPI name is `probnode-sdk` because `probnode` is taken.
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
# probnode-sdk
|
|
2
|
+
|
|
3
|
+
Python client and command line for [ProbNode](https://probnode.com): Polymarket order books, trades
|
|
4
|
+
and feed gaps as hourly and market-day files (lp2, parquet, csv.gz), with a sha256 for every file.
|
|
5
|
+
|
|
6
|
+
```sh
|
|
7
|
+
pip install ./packages/probnode # from a checkout; on PyPI as probnode-sdk soon
|
|
8
|
+
# (not `pip install probnode`: that is someone else's package)
|
|
9
|
+
probnode login
|
|
10
|
+
probnode datasets
|
|
11
|
+
probnode files books parquet 2026-11-08T12 2026-11-08T14
|
|
12
|
+
probnode sync books parquet 2026-11-08T12 2026-11-08T14 -o data/ # again later: only changes
|
|
13
|
+
probnode doctor data/ # re-check every sha256
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
import probnode
|
|
18
|
+
|
|
19
|
+
pn = probnode.Client() # PROBNODE_API_KEY, or the key `probnode login` saved
|
|
20
|
+
files = pn.files("books", "parquet", "2026-11-08T12:00:00Z", "2026-11-08T13:00:00Z")
|
|
21
|
+
paths = list(pn.download_all(files, "data/")) # each checked against its sha256
|
|
22
|
+
table = probnode.read(paths[0]) # pyarrow.Table -> .to_pandas() / polars.from_arrow()
|
|
23
|
+
|
|
24
|
+
for batch in probnode.iter_lp2("hour.lp2"): # native books, decoded in bounded memory
|
|
25
|
+
...
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
Units: prices x10,000 (5400 = 0.54), sizes x1,000,000 shares, timestamps in ms UTC. Signed links
|
|
29
|
+
last 15 minutes; the client signs again whenever it needs to. The import name is `probnode`; the
|
|
30
|
+
PyPI name is `probnode-sdk` because `probnode` is taken.
|
|
@@ -0,0 +1,43 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "probnode-sdk"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Python client and CLI for ProbNode: Polymarket order books, trades and gaps"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10"
|
|
7
|
+
license = { text = "MIT" }
|
|
8
|
+
dependencies = [
|
|
9
|
+
"httpx>=0.27",
|
|
10
|
+
# The lp2 reader decodes the native book format; parquet and csv.gz need nothing extra to
|
|
11
|
+
# download, but reading any of them into a table does.
|
|
12
|
+
"numpy>=1.24",
|
|
13
|
+
"pyarrow>=14",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[project.optional-dependencies]
|
|
17
|
+
pandas = ["pandas>=2.0"]
|
|
18
|
+
polars = ["polars>=0.20"]
|
|
19
|
+
dev = ["pytest>=8.3", "ruff>=0.6"]
|
|
20
|
+
|
|
21
|
+
[project.urls]
|
|
22
|
+
Homepage = "https://probnode.com"
|
|
23
|
+
Documentation = "https://probnode.com/docs"
|
|
24
|
+
|
|
25
|
+
[project.scripts]
|
|
26
|
+
probnode = "probnode.cli:main"
|
|
27
|
+
|
|
28
|
+
[build-system]
|
|
29
|
+
requires = ["hatchling"]
|
|
30
|
+
build-backend = "hatchling.build"
|
|
31
|
+
|
|
32
|
+
[tool.hatch.build.targets.wheel]
|
|
33
|
+
packages = ["src/probnode"]
|
|
34
|
+
|
|
35
|
+
[tool.pytest.ini_options]
|
|
36
|
+
testpaths = ["tests"]
|
|
37
|
+
|
|
38
|
+
[tool.ruff]
|
|
39
|
+
line-length = 100
|
|
40
|
+
target-version = "py310"
|
|
41
|
+
|
|
42
|
+
[tool.ruff.lint]
|
|
43
|
+
select = ["E", "F", "I", "UP", "B"]
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
"""ProbNode: Polymarket order books, trades and feed gaps, as files and as a Python API."""
|
|
2
|
+
|
|
3
|
+
from .client import Client, DataFile, ProbNodeError
|
|
4
|
+
from .read import iter_lp2, read
|
|
5
|
+
|
|
6
|
+
__all__ = ["Client", "DataFile", "ProbNodeError", "iter_lp2", "read"]
|
|
7
|
+
__version__ = "0.1.0"
|
|
@@ -0,0 +1,227 @@
|
|
|
1
|
+
"""`probnode` on the command line.
|
|
2
|
+
|
|
3
|
+
probnode login # email + password -> a key in ~/.config
|
|
4
|
+
probnode datasets
|
|
5
|
+
probnode files books parquet 2026-11-08T12 2026-11-08T14
|
|
6
|
+
probnode download books parquet 2026-11-08T12 2026-11-08T14 -o data/
|
|
7
|
+
probnode sync books parquet 2026-11-08T12 2026-11-08T14 -o data/ # only what changed
|
|
8
|
+
probnode doctor data/ # every file against its sha256
|
|
9
|
+
|
|
10
|
+
Times are UTC hours: `2026-11-08T12`, `2026-11-08T12:00`, or full ISO all work.
|
|
11
|
+
Exit codes: 0 done, 1 a check failed (doctor), 2 the API refused or a request failed.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import argparse
|
|
17
|
+
import getpass
|
|
18
|
+
import json
|
|
19
|
+
import sys
|
|
20
|
+
from pathlib import Path
|
|
21
|
+
|
|
22
|
+
from .client import Client, DataFile, ProbNodeError, load_config, save_config, sha256_of
|
|
23
|
+
|
|
24
|
+
STATE = ".probnode-sync.json"
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def iso_hour(value: str) -> str:
|
|
28
|
+
"""`2026-11-08T12` / `2026-11-08 12:00` / ISO -> `2026-11-08T12:00:00Z`."""
|
|
29
|
+
v = value.strip().replace(" ", "T").rstrip("Z")
|
|
30
|
+
date, _, time = v.partition("T")
|
|
31
|
+
hour = (time or "00")[:2]
|
|
32
|
+
if len(date) != 10 or not hour.isdigit():
|
|
33
|
+
raise argparse.ArgumentTypeError(f"not a UTC hour: {value!r} (use 2026-11-08T12)")
|
|
34
|
+
return f"{date}T{hour}:00:00Z"
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def human(n: int) -> str:
|
|
38
|
+
size = float(n)
|
|
39
|
+
for unit in ("B", "KB", "MB", "GB"):
|
|
40
|
+
if size < 1024 or unit == "GB":
|
|
41
|
+
return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} B"
|
|
42
|
+
size /= 1024
|
|
43
|
+
return f"{n} B"
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _client(args: argparse.Namespace) -> Client:
|
|
47
|
+
return Client(api_key=args.key, base_url=args.base_url)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def _listing(pn: Client, args: argparse.Namespace) -> list[DataFile]:
|
|
51
|
+
return pn.files(
|
|
52
|
+
args.dataset,
|
|
53
|
+
args.format,
|
|
54
|
+
args.start,
|
|
55
|
+
args.end,
|
|
56
|
+
granularity=args.granularity,
|
|
57
|
+
market=args.market or (),
|
|
58
|
+
series=args.series or (),
|
|
59
|
+
)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _read_state(dest: Path) -> dict[str, dict]:
|
|
63
|
+
try:
|
|
64
|
+
return json.loads((dest / STATE).read_text(encoding="utf-8"))
|
|
65
|
+
except FileNotFoundError:
|
|
66
|
+
return {}
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _write_state(dest: Path, state: dict[str, dict]) -> None:
|
|
70
|
+
tmp = dest / (STATE + ".tmp")
|
|
71
|
+
tmp.write_text(json.dumps(state, indent=1, sort_keys=True), encoding="utf-8")
|
|
72
|
+
tmp.replace(dest / STATE)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def _current(dest: Path, recorded: dict | None, f: DataFile) -> bool:
|
|
76
|
+
"""A file is current when the listing still names the sha256 we fetched AND the file on disk
|
|
77
|
+
is the one we wrote: same size, same mtime. Anything touched since - truncated, overwritten,
|
|
78
|
+
restored from elsewhere - is fetched again, rather than trusted because the state says so.
|
|
79
|
+
(A full re-hash is `probnode doctor`.)"""
|
|
80
|
+
path = dest / f.object_path
|
|
81
|
+
if not recorded or recorded.get("sha256") != f.sha256 or not path.exists():
|
|
82
|
+
return False
|
|
83
|
+
st = path.stat()
|
|
84
|
+
return st.st_size == f.bytes and st.st_mtime_ns == recorded.get("mtime_ns")
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def cmd_login(args: argparse.Namespace) -> int:
|
|
88
|
+
email = args.email or input("Email: ")
|
|
89
|
+
password = getpass.getpass("Password: ")
|
|
90
|
+
pn = _client(args)
|
|
91
|
+
key = pn.login(email, password)
|
|
92
|
+
config = load_config()
|
|
93
|
+
config.update({"api_key": key, "base_url": pn.base_url})
|
|
94
|
+
save_config(config)
|
|
95
|
+
print(f"signed in as {email}; key saved")
|
|
96
|
+
return 0
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def cmd_datasets(args: argparse.Namespace) -> int:
|
|
100
|
+
for d in _client(args).datasets():
|
|
101
|
+
held = f"{d['first_hour']} .. {d['last_hour']}" if d.get("last_hour") else "nothing held"
|
|
102
|
+
print(f"{d['code']:<10} {d['tier']:<5} {','.join(d['formats']):<20} {held}")
|
|
103
|
+
return 0
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def cmd_files(args: argparse.Namespace) -> int:
|
|
107
|
+
files = _listing(_client(args), args)
|
|
108
|
+
for f in files:
|
|
109
|
+
print(f"{f.object_path} {f.rows:>12,} rows {human(f.bytes):>10} {f.sha256[:12]}")
|
|
110
|
+
print(
|
|
111
|
+
f"{len(files)} files, {human(sum(f.bytes for f in files))}, "
|
|
112
|
+
f"{sum(f.rows for f in files):,} rows"
|
|
113
|
+
)
|
|
114
|
+
return 0
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def cmd_download(args: argparse.Namespace) -> int:
|
|
118
|
+
pn = _client(args)
|
|
119
|
+
files = _listing(pn, args)
|
|
120
|
+
dest = Path(args.out)
|
|
121
|
+
for path in pn.download_all(files, dest):
|
|
122
|
+
print(path)
|
|
123
|
+
return 0
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def cmd_sync(args: argparse.Namespace) -> int:
|
|
127
|
+
"""Download what is new or changed since the last sync into the same directory; a file whose
|
|
128
|
+
sha256 already matches is left alone, so a second run moves nothing. Interrupted downloads
|
|
129
|
+
leave only a .part file, and the next run fetches that file again."""
|
|
130
|
+
pn = _client(args)
|
|
131
|
+
dest = Path(args.out)
|
|
132
|
+
dest.mkdir(parents=True, exist_ok=True)
|
|
133
|
+
state = _read_state(dest)
|
|
134
|
+
files = _listing(pn, args)
|
|
135
|
+
todo = [f for f in files if not _current(dest, state.get(f.object_path), f)]
|
|
136
|
+
moved = 0
|
|
137
|
+
for f, path in zip(todo, pn.download_all(todo, dest), strict=False):
|
|
138
|
+
st = path.stat()
|
|
139
|
+
state[f.object_path] = {
|
|
140
|
+
"sha256": f.sha256,
|
|
141
|
+
"bytes": f.bytes,
|
|
142
|
+
"rows": f.rows,
|
|
143
|
+
"mtime_ns": st.st_mtime_ns,
|
|
144
|
+
}
|
|
145
|
+
_write_state(dest, state)
|
|
146
|
+
moved += f.bytes
|
|
147
|
+
print(f"got {path}")
|
|
148
|
+
print(
|
|
149
|
+
f"{len(todo)} of {len(files)} files fetched, {human(moved)}; "
|
|
150
|
+
f"{len(files) - len(todo)} already current"
|
|
151
|
+
)
|
|
152
|
+
return 0
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def cmd_doctor(args: argparse.Namespace) -> int:
|
|
156
|
+
"""Every file a sync recorded, re-hashed. Missing or changed files are named; exit 1 if any."""
|
|
157
|
+
dest = Path(args.dir)
|
|
158
|
+
state = _read_state(dest)
|
|
159
|
+
if not state:
|
|
160
|
+
print(f"no {STATE} in {dest}: run `probnode sync` into it first", file=sys.stderr)
|
|
161
|
+
return 1
|
|
162
|
+
bad = 0
|
|
163
|
+
for object_path, meta in sorted(state.items()):
|
|
164
|
+
path = dest / object_path
|
|
165
|
+
if not path.exists():
|
|
166
|
+
print(f"MISSING {object_path}")
|
|
167
|
+
bad += 1
|
|
168
|
+
elif sha256_of(path) != meta["sha256"]:
|
|
169
|
+
print(f"CORRUPT {object_path}")
|
|
170
|
+
bad += 1
|
|
171
|
+
print(f"{len(state) - bad} of {len(state)} files ok")
|
|
172
|
+
return 1 if bad else 0
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
def parser() -> argparse.ArgumentParser:
|
|
176
|
+
p = argparse.ArgumentParser(prog="probnode", description="ProbNode data from the shell")
|
|
177
|
+
p.add_argument("--key", help="API key (default: PROBNODE_API_KEY or `probnode login`)")
|
|
178
|
+
p.add_argument("--base-url", help="API base (default: https://probnode.com)")
|
|
179
|
+
sub = p.add_subparsers(dest="command", required=True)
|
|
180
|
+
|
|
181
|
+
login = sub.add_parser("login", help="sign in with email and password, save the key")
|
|
182
|
+
login.add_argument("--email")
|
|
183
|
+
login.set_defaults(run=cmd_login)
|
|
184
|
+
|
|
185
|
+
sub.add_parser("datasets", help="datasets, formats and the hours held").set_defaults(
|
|
186
|
+
run=cmd_datasets
|
|
187
|
+
)
|
|
188
|
+
|
|
189
|
+
def selection(sp: argparse.ArgumentParser) -> None:
|
|
190
|
+
sp.add_argument("dataset", help="books, trades, gaps, markets, quotes, ohlcv_1m")
|
|
191
|
+
sp.add_argument("format", help="lp2, parquet or csv")
|
|
192
|
+
sp.add_argument("start", type=iso_hour, help="first UTC hour, inclusive")
|
|
193
|
+
sp.add_argument("end", type=iso_hour, help="last UTC hour, exclusive")
|
|
194
|
+
sp.add_argument("--granularity", default="hour", choices=["hour", "market_day"])
|
|
195
|
+
sp.add_argument("--market", action="append", help="narrow to a market slug (repeatable)")
|
|
196
|
+
sp.add_argument("--series", action="append", help="narrow to a series slug (repeatable)")
|
|
197
|
+
|
|
198
|
+
files = sub.add_parser("files", help="list the files in a selection")
|
|
199
|
+
selection(files)
|
|
200
|
+
files.set_defaults(run=cmd_files)
|
|
201
|
+
|
|
202
|
+
for name, run, text in (
|
|
203
|
+
("download", cmd_download, "download a selection, sha256-checked"),
|
|
204
|
+
("sync", cmd_sync, "download only what is new or changed"),
|
|
205
|
+
):
|
|
206
|
+
sp = sub.add_parser(name, help=text)
|
|
207
|
+
selection(sp)
|
|
208
|
+
sp.add_argument("-o", "--out", default=".", help="directory (default: here)")
|
|
209
|
+
sp.set_defaults(run=run)
|
|
210
|
+
|
|
211
|
+
doctor = sub.add_parser("doctor", help="re-check every synced file against its sha256")
|
|
212
|
+
doctor.add_argument("dir", nargs="?", default=".")
|
|
213
|
+
doctor.set_defaults(run=cmd_doctor)
|
|
214
|
+
return p
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def main(argv: list[str] | None = None) -> int:
|
|
218
|
+
args = parser().parse_args(argv)
|
|
219
|
+
try:
|
|
220
|
+
return args.run(args)
|
|
221
|
+
except ProbNodeError as exc:
|
|
222
|
+
print(f"probnode: {exc}", file=sys.stderr)
|
|
223
|
+
return 2
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
if __name__ == "__main__":
|
|
227
|
+
sys.exit(main())
|
|
@@ -0,0 +1,289 @@
|
|
|
1
|
+
"""The ProbNode API, as a Python object.
|
|
2
|
+
|
|
3
|
+
import probnode
|
|
4
|
+
pn = probnode.Client() # key from PROBNODE_API_KEY or `probnode login`
|
|
5
|
+
files = pn.files("books", "parquet", "2026-11-08T12:00:00Z", "2026-11-08T14:00:00Z")
|
|
6
|
+
for f in files:
|
|
7
|
+
path = pn.download(f, "data/") # sha256-checked against the listing
|
|
8
|
+
table = probnode.read(path) # pyarrow.Table, whatever the format
|
|
9
|
+
|
|
10
|
+
Every method is one HTTP call (download is two: sign, then fetch). Errors are `ProbNodeError` with
|
|
11
|
+
the API's own sentence in it, never a bare status code.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import hashlib
|
|
17
|
+
import json
|
|
18
|
+
import os
|
|
19
|
+
from collections.abc import Iterable, Iterator
|
|
20
|
+
from dataclasses import dataclass, field
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
from typing import Any
|
|
23
|
+
|
|
24
|
+
import httpx
|
|
25
|
+
|
|
26
|
+
DEFAULT_BASE_URL = "https://probnode.com"
|
|
27
|
+
CONFIG_PATH = Path(os.environ.get("PROBNODE_CONFIG", Path.home() / ".config" / "probnode.json"))
|
|
28
|
+
SIGN_BATCH = 100 # POST /v1/data/urls takes at most 100 paths
|
|
29
|
+
CHUNK = 1 << 20
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class ProbNodeError(Exception):
|
|
33
|
+
"""A refusal or failure from the API, with its status and its sentence."""
|
|
34
|
+
|
|
35
|
+
def __init__(self, status: int, detail: str):
|
|
36
|
+
super().__init__(f"{status}: {detail}")
|
|
37
|
+
self.status = status
|
|
38
|
+
self.detail = detail
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
@dataclass(frozen=True)
|
|
42
|
+
class DataFile:
|
|
43
|
+
"""One object in the lake, as GET /v1/data/files lists it."""
|
|
44
|
+
|
|
45
|
+
dataset: str
|
|
46
|
+
format: str
|
|
47
|
+
event_date: str
|
|
48
|
+
hour: int
|
|
49
|
+
src: str
|
|
50
|
+
object_path: str
|
|
51
|
+
rows: int
|
|
52
|
+
bytes: int
|
|
53
|
+
sha256: str
|
|
54
|
+
granularity: str = "hour"
|
|
55
|
+
market: str | None = None
|
|
56
|
+
extra: dict[str, Any] = field(default_factory=dict, compare=False, repr=False)
|
|
57
|
+
|
|
58
|
+
@classmethod
|
|
59
|
+
def from_api(cls, raw: dict[str, Any]) -> DataFile:
|
|
60
|
+
known = {k: raw[k] for k in cls.__dataclass_fields__ if k in raw and k != "extra"}
|
|
61
|
+
rest = {k: v for k, v in raw.items() if k not in cls.__dataclass_fields__}
|
|
62
|
+
return cls(**known, extra=rest)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def load_config() -> dict[str, str]:
|
|
66
|
+
try:
|
|
67
|
+
return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
|
|
68
|
+
except FileNotFoundError:
|
|
69
|
+
return {}
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def save_config(values: dict[str, str]) -> None:
|
|
73
|
+
CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
74
|
+
CONFIG_PATH.write_text(json.dumps(values, indent=2), encoding="utf-8")
|
|
75
|
+
try:
|
|
76
|
+
CONFIG_PATH.chmod(0o600)
|
|
77
|
+
except OSError:
|
|
78
|
+
pass # Windows: the file lives in the user's profile, which is already private
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def sha256_of(path: Path) -> str:
|
|
82
|
+
h = hashlib.sha256()
|
|
83
|
+
with path.open("rb") as f:
|
|
84
|
+
for chunk in iter(lambda: f.read(CHUNK), b""):
|
|
85
|
+
h.update(chunk)
|
|
86
|
+
return h.hexdigest()
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
class Client:
|
|
90
|
+
"""A session against one ProbNode deployment.
|
|
91
|
+
|
|
92
|
+
`api_key` and `base_url` fall back to PROBNODE_API_KEY / PROBNODE_BASE_URL, then to what
|
|
93
|
+
`probnode login` saved, then (base URL only) to https://probnode.com. Public calls - datasets,
|
|
94
|
+
samples, the catalog - work without a key.
|
|
95
|
+
"""
|
|
96
|
+
|
|
97
|
+
def __init__(
|
|
98
|
+
self,
|
|
99
|
+
api_key: str | None = None,
|
|
100
|
+
base_url: str | None = None,
|
|
101
|
+
*,
|
|
102
|
+
timeout: float = 60.0,
|
|
103
|
+
transport: httpx.BaseTransport | None = None,
|
|
104
|
+
):
|
|
105
|
+
saved = load_config()
|
|
106
|
+
self.api_key = api_key or os.environ.get("PROBNODE_API_KEY") or saved.get("api_key")
|
|
107
|
+
self.base_url = (
|
|
108
|
+
base_url
|
|
109
|
+
or os.environ.get("PROBNODE_BASE_URL")
|
|
110
|
+
or saved.get("base_url")
|
|
111
|
+
or DEFAULT_BASE_URL
|
|
112
|
+
).rstrip("/")
|
|
113
|
+
self._http = httpx.Client(
|
|
114
|
+
base_url=self.base_url, timeout=timeout, transport=transport, follow_redirects=False
|
|
115
|
+
)
|
|
116
|
+
# Signed URLs point at the lake, possibly another host; fetched without our key.
|
|
117
|
+
self._fetch = httpx.Client(timeout=timeout, transport=transport, follow_redirects=True)
|
|
118
|
+
|
|
119
|
+
def __enter__(self) -> Client:
|
|
120
|
+
return self
|
|
121
|
+
|
|
122
|
+
def __exit__(self, *exc: object) -> None:
|
|
123
|
+
self.close()
|
|
124
|
+
|
|
125
|
+
def close(self) -> None:
|
|
126
|
+
self._http.close()
|
|
127
|
+
self._fetch.close()
|
|
128
|
+
|
|
129
|
+
# -- plumbing ---------------------------------------------------------------------------
|
|
130
|
+
|
|
131
|
+
def _headers(self, need_key: bool) -> dict[str, str]:
|
|
132
|
+
if self.api_key:
|
|
133
|
+
return {"Authorization": f"Bearer {self.api_key}"}
|
|
134
|
+
if need_key:
|
|
135
|
+
raise ProbNodeError(
|
|
136
|
+
401, "no API key: pass api_key, set PROBNODE_API_KEY, or run `probnode login`"
|
|
137
|
+
)
|
|
138
|
+
return {}
|
|
139
|
+
|
|
140
|
+
def _request(self, method: str, path: str, *, need_key: bool = True, **kw: Any) -> Any:
|
|
141
|
+
response = self._http.request(method, path, headers=self._headers(need_key), **kw)
|
|
142
|
+
if response.status_code >= 400:
|
|
143
|
+
try:
|
|
144
|
+
detail = response.json().get("detail")
|
|
145
|
+
except ValueError:
|
|
146
|
+
detail = None
|
|
147
|
+
if not isinstance(detail, str):
|
|
148
|
+
detail = json.dumps(detail) if detail else response.reason_phrase
|
|
149
|
+
raise ProbNodeError(response.status_code, detail)
|
|
150
|
+
return response
|
|
151
|
+
|
|
152
|
+
def _json(self, method: str, path: str, **kw: Any) -> Any:
|
|
153
|
+
return self._request(method, path, **kw).json()
|
|
154
|
+
|
|
155
|
+
# -- account ----------------------------------------------------------------------------
|
|
156
|
+
|
|
157
|
+
def login(self, email: str, password: str) -> str:
|
|
158
|
+
"""Mint a key with email and password; the client uses it from then on."""
|
|
159
|
+
out = self._json(
|
|
160
|
+
"POST", "/v1/auth/login", need_key=False, json={"email": email, "password": password}
|
|
161
|
+
)
|
|
162
|
+
self.api_key = out["key"]
|
|
163
|
+
return self.api_key
|
|
164
|
+
|
|
165
|
+
def me(self) -> dict[str, Any]:
|
|
166
|
+
return self._json("GET", "/v1/me")
|
|
167
|
+
|
|
168
|
+
def usage(self) -> dict[str, Any]:
|
|
169
|
+
return self._json("GET", "/v1/data/usage")
|
|
170
|
+
|
|
171
|
+
# -- catalog ----------------------------------------------------------------------------
|
|
172
|
+
|
|
173
|
+
def datasets(self) -> list[dict[str, Any]]:
|
|
174
|
+
"""Every dataset, its formats, its tier and the hours held. Public."""
|
|
175
|
+
return self._json("GET", "/v1/data/datasets", need_key=False)
|
|
176
|
+
|
|
177
|
+
def series(self, category: str | None = None, q: str | None = None) -> list[dict[str, Any]]:
|
|
178
|
+
params = {k: v for k, v in {"category": category, "q": q}.items() if v}
|
|
179
|
+
return self._json("GET", "/v1/catalog/series", need_key=False, params=params)
|
|
180
|
+
|
|
181
|
+
def markets(self, series: str | None = None, **params: Any) -> Any:
|
|
182
|
+
if series:
|
|
183
|
+
params["series"] = series
|
|
184
|
+
return self._json("GET", "/v1/catalog/markets", need_key=False, params=params)
|
|
185
|
+
|
|
186
|
+
def coverage(self, token_id: str, start: str | None = None, end: str | None = None) -> Any:
|
|
187
|
+
params = {k: v for k, v in {"token_id": token_id, "start": start, "end": end}.items() if v}
|
|
188
|
+
return self._json("GET", "/v1/catalog/coverage", need_key=False, params=params)
|
|
189
|
+
|
|
190
|
+
def samples(self) -> Any:
|
|
191
|
+
"""Free sample files, one market-day per dataset and format. Public."""
|
|
192
|
+
return self._json("GET", "/v1/samples", need_key=False)
|
|
193
|
+
|
|
194
|
+
# -- files ------------------------------------------------------------------------------
|
|
195
|
+
|
|
196
|
+
def files(
|
|
197
|
+
self,
|
|
198
|
+
dataset: str,
|
|
199
|
+
format: str,
|
|
200
|
+
start: str,
|
|
201
|
+
end: str,
|
|
202
|
+
*,
|
|
203
|
+
granularity: str = "hour",
|
|
204
|
+
market: Iterable[str] = (),
|
|
205
|
+
series: Iterable[str] = (),
|
|
206
|
+
) -> list[DataFile]:
|
|
207
|
+
"""The objects for one dataset and format over [start, end), from the manifests."""
|
|
208
|
+
params: list[tuple[str, str]] = [
|
|
209
|
+
("dataset", dataset),
|
|
210
|
+
("format", format),
|
|
211
|
+
("start", start),
|
|
212
|
+
("end", end),
|
|
213
|
+
]
|
|
214
|
+
if granularity != "hour":
|
|
215
|
+
params.append(("granularity", granularity))
|
|
216
|
+
params += [("market", m) for m in market] + [("series", s) for s in series]
|
|
217
|
+
out = self._json("GET", "/v1/data/files", params=params)
|
|
218
|
+
return [DataFile.from_api(f) for f in out["files"]]
|
|
219
|
+
|
|
220
|
+
def sign(self, paths: Iterable[str]) -> dict[str, str]:
|
|
221
|
+
"""object_path -> signed URL (valid for minutes; sign again any time)."""
|
|
222
|
+
paths = list(paths)
|
|
223
|
+
urls: dict[str, str] = {}
|
|
224
|
+
for i in range(0, len(paths), SIGN_BATCH):
|
|
225
|
+
out = self._json("POST", "/v1/data/urls", json={"paths": paths[i : i + SIGN_BATCH]})
|
|
226
|
+
urls.update({u["object_path"]: u["url"] for u in out["urls"]})
|
|
227
|
+
return urls
|
|
228
|
+
|
|
229
|
+
def download(
|
|
230
|
+
self,
|
|
231
|
+
file: DataFile,
|
|
232
|
+
dest: str | Path = ".",
|
|
233
|
+
*,
|
|
234
|
+
url: str | None = None,
|
|
235
|
+
verify: bool = True,
|
|
236
|
+
) -> Path:
|
|
237
|
+
"""Fetch one file to `dest/<object_path>`, written to a .part first and renamed only when
|
|
238
|
+
complete and (unless verify=False) its sha256 matches the listing."""
|
|
239
|
+
target = Path(dest) / file.object_path
|
|
240
|
+
target.parent.mkdir(parents=True, exist_ok=True)
|
|
241
|
+
url = url or self.sign([file.object_path])[file.object_path]
|
|
242
|
+
part = target.with_name(target.name + ".part")
|
|
243
|
+
h = hashlib.sha256()
|
|
244
|
+
with self._fetch.stream("GET", url) as response:
|
|
245
|
+
if response.status_code >= 400:
|
|
246
|
+
raise ProbNodeError(response.status_code, f"fetching {file.object_path} failed")
|
|
247
|
+
with part.open("wb") as out:
|
|
248
|
+
for chunk in response.iter_bytes(CHUNK):
|
|
249
|
+
out.write(chunk)
|
|
250
|
+
h.update(chunk)
|
|
251
|
+
if verify and h.hexdigest() != file.sha256:
|
|
252
|
+
part.unlink(missing_ok=True)
|
|
253
|
+
raise ProbNodeError(0, f"sha256 mismatch for {file.object_path}")
|
|
254
|
+
part.replace(target)
|
|
255
|
+
return target
|
|
256
|
+
|
|
257
|
+
def download_all(
|
|
258
|
+
self, files: Iterable[DataFile], dest: str | Path = ".", *, verify: bool = True
|
|
259
|
+
) -> Iterator[Path]:
|
|
260
|
+
files = list(files)
|
|
261
|
+
urls = self.sign(f.object_path for f in files)
|
|
262
|
+
for f in files:
|
|
263
|
+
yield self.download(f, dest, url=urls[f.object_path], verify=verify)
|
|
264
|
+
|
|
265
|
+
# -- range query ------------------------------------------------------------------------
|
|
266
|
+
|
|
267
|
+
def timeseries(
|
|
268
|
+
self, dataset: str, market: str, start: str, end: str, *, format: str = "parquet"
|
|
269
|
+
) -> Any:
|
|
270
|
+
"""Rows for one market over a window (at most 7 days, 200,000 rows) as a pyarrow Table
|
|
271
|
+
(format="parquet") or a list of dicts (format="ndjson")."""
|
|
272
|
+
response = self._request(
|
|
273
|
+
"GET",
|
|
274
|
+
"/v1/timeseries",
|
|
275
|
+
params={
|
|
276
|
+
"dataset": dataset,
|
|
277
|
+
"market": market,
|
|
278
|
+
"start": start,
|
|
279
|
+
"end": end,
|
|
280
|
+
"format": format,
|
|
281
|
+
},
|
|
282
|
+
)
|
|
283
|
+
if format == "ndjson":
|
|
284
|
+
return [json.loads(line) for line in response.text.splitlines() if line]
|
|
285
|
+
import io
|
|
286
|
+
|
|
287
|
+
import pyarrow.parquet as pq
|
|
288
|
+
|
|
289
|
+
return pq.read_table(io.BytesIO(response.content))
|
|
@@ -0,0 +1,357 @@
|
|
|
1
|
+
"""The recorder's LP2 book archive, decoded. Vendored, not imported.
|
|
2
|
+
|
|
3
|
+
The reference reader for the LP2 format ProbNode's recorder writes, kept byte-for-byte equal to
|
|
4
|
+
the copy the ProbNode ingest uses, so what a customer decodes is what the product exported.
|
|
5
|
+
|
|
6
|
+
What changed against the original, and why:
|
|
7
|
+
|
|
8
|
+
* `decode` returns the per-level columns (`bid_px0`, `bid_sz0`, ...) the decoder already holds,
|
|
9
|
+
instead of only folding them back into list columns. The exports want exactly those columns
|
|
10
|
+
(the CSV is flat), and rebuilding lists only to flatten them again costs a quarter of the time.
|
|
11
|
+
* `to_table` is lobpack's `from_arrays`, for callers that want the books schema as the recorder
|
|
12
|
+
defines it - the round-trip tests compare against it.
|
|
13
|
+
* `stream` is new: the same decoding a batch of whole markets at a time, in bounded memory, for
|
|
14
|
+
the ingest. `decode` stays as the reference the tests hold `stream` to.
|
|
15
|
+
* LP1 (the xz prototype) is not here. The lake holds LP2 only.
|
|
16
|
+
|
|
17
|
+
The format, in one paragraph (the Go file is the authority): `b"LOBP2\\n"`, a u32 LE header
|
|
18
|
+
length, a JSON header, then one zstd frame per stream in header order. Rows are grouped per token
|
|
19
|
+
and stored change-only - every field a zigzag delta against the previous row of the same token,
|
|
20
|
+
written only when it changed, with a per-row bit mask saying which. The NO book of a market is the
|
|
21
|
+
exact complement of the YES book, so a row of token 2k+1 that mirrors the row of token 2k at the
|
|
22
|
+
same millisecond is not stored at all: its partner carries a `twin` bit. Prices are x10000, sizes
|
|
23
|
+
x1e6, levels best first.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
import json
|
|
27
|
+
import struct
|
|
28
|
+
from collections.abc import Iterator
|
|
29
|
+
from dataclasses import dataclass
|
|
30
|
+
from typing import Any
|
|
31
|
+
|
|
32
|
+
import numpy as np
|
|
33
|
+
import pyarrow as pa
|
|
34
|
+
|
|
35
|
+
MAGIC2 = b"LOBP2\n"
|
|
36
|
+
VERSION = 2
|
|
37
|
+
PX_ONE = 10000 # px ticks of 1e-4: price 1.0
|
|
38
|
+
|
|
39
|
+
#: The books schema exactly as lobpack defines it, so `to_table` output compares equal to the
|
|
40
|
+
#: reference reader's.
|
|
41
|
+
SCHEMA = pa.schema(
|
|
42
|
+
[
|
|
43
|
+
("ts_ms", pa.int64(), False),
|
|
44
|
+
("ts_ex_ms", pa.int64(), False),
|
|
45
|
+
("tok", pa.int32(), False),
|
|
46
|
+
("bid_px", pa.list_(pa.field("element", pa.int32(), False)), False),
|
|
47
|
+
("bid_sz", pa.list_(pa.field("element", pa.int64(), False)), False),
|
|
48
|
+
("ask_px", pa.list_(pa.field("element", pa.int32(), False)), False),
|
|
49
|
+
("ask_sz", pa.list_(pa.field("element", pa.int64(), False)), False),
|
|
50
|
+
("snap", pa.bool_(), False),
|
|
51
|
+
("shard", pa.int32(), False),
|
|
52
|
+
("conn", pa.int32(), False),
|
|
53
|
+
]
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
SIDES = ("bid_px", "bid_sz", "ask_px", "ask_sz")
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
@dataclass(frozen=True)
|
|
60
|
+
class Decoded:
|
|
61
|
+
"""One LP2 object as columns, in (ts_ms, tok) order like the reference reader.
|
|
62
|
+
|
|
63
|
+
`columns` holds `ts_ms`, `ts_ex_ms`, `tok`, `snap`, `shard`, `conn`, the level counts `nb` and
|
|
64
|
+
`na`, and one array per side and level (`bid_px0` ... `ask_sz{depth-1}`). An absent level is 0
|
|
65
|
+
in its array and is outside `nb`/`na`; nothing else marks it.
|
|
66
|
+
"""
|
|
67
|
+
|
|
68
|
+
header: dict[str, Any]
|
|
69
|
+
depth: int
|
|
70
|
+
columns: dict[str, np.ndarray]
|
|
71
|
+
|
|
72
|
+
@property
|
|
73
|
+
def rows(self) -> int:
|
|
74
|
+
return int(self.columns["ts_ms"].size)
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
# ---------- varint / zigzag (vectorised), as in lobpack ----------
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def unzz(u: np.ndarray) -> np.ndarray:
|
|
81
|
+
u = u.astype(np.uint64)
|
|
82
|
+
return ((u >> np.uint64(1)).astype(np.int64)) ^ (-(u & np.uint64(1)).astype(np.int64))
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def vdec(buf: bytes, count: int) -> np.ndarray:
|
|
86
|
+
b = np.frombuffer(buf, np.uint8)
|
|
87
|
+
ends = np.flatnonzero(b < 0x80)
|
|
88
|
+
if ends.size != count:
|
|
89
|
+
raise ValueError(f"varint stream has {ends.size} values, header says {count}")
|
|
90
|
+
if count == 0:
|
|
91
|
+
return np.zeros(0, np.uint64)
|
|
92
|
+
starts = np.r_[0, ends[:-1] + 1]
|
|
93
|
+
lens = ends - starts + 1
|
|
94
|
+
out = np.zeros(count, np.uint64)
|
|
95
|
+
for k in range(int(lens.max())):
|
|
96
|
+
m = lens > k
|
|
97
|
+
out[m] |= (b[starts[m] + k] & np.uint64(0x7F)).astype(np.uint64) << np.uint64(7 * k)
|
|
98
|
+
return out
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def seg_cumsum(d: np.ndarray, run_start: np.ndarray, run_len: np.ndarray) -> np.ndarray:
|
|
102
|
+
c = np.cumsum(d)
|
|
103
|
+
before = np.r_[0, c[:-1]][run_start]
|
|
104
|
+
return c - np.repeat(before, run_len)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def mirror_of(F: dict[str, np.ndarray], depth: int, i: np.ndarray) -> dict[str, np.ndarray]:
|
|
108
|
+
"""Complement rows of rows i (token 2k -> 2k+1)."""
|
|
109
|
+
G = {
|
|
110
|
+
"ts_ms": F["ts_ms"][i],
|
|
111
|
+
"ts_ex_ms": F["ts_ex_ms"][i],
|
|
112
|
+
"tok": F["tok"][i] + 1,
|
|
113
|
+
"snap": F["snap"][i],
|
|
114
|
+
"shard": F["shard"][i],
|
|
115
|
+
"conn": F["conn"][i],
|
|
116
|
+
"nb": F["na"][i],
|
|
117
|
+
"na": F["nb"][i],
|
|
118
|
+
}
|
|
119
|
+
for k in range(depth):
|
|
120
|
+
G[f"bid_px{k}"] = np.where(F["na"][i] > k, PX_ONE - F[f"ask_px{k}"][i], 0)
|
|
121
|
+
G[f"bid_sz{k}"] = np.where(F["na"][i] > k, F[f"ask_sz{k}"][i], 0)
|
|
122
|
+
G[f"ask_px{k}"] = np.where(F["nb"][i] > k, PX_ONE - F[f"bid_px{k}"][i], 0)
|
|
123
|
+
G[f"ask_sz{k}"] = np.where(F["nb"][i] > k, F[f"bid_sz{k}"][i], 0)
|
|
124
|
+
return G
|
|
125
|
+
|
|
126
|
+
|
|
127
|
+
def read_header(blob: bytes) -> tuple[dict[str, Any], int]:
|
|
128
|
+
"""The JSON header and the offset of the first stream. Cheap: reads no stream."""
|
|
129
|
+
if blob[: len(MAGIC2)] != MAGIC2:
|
|
130
|
+
raise ValueError("not an LP2 object")
|
|
131
|
+
p = len(MAGIC2)
|
|
132
|
+
(hl,) = struct.unpack("<I", blob[p : p + 4])
|
|
133
|
+
hdr = json.loads(blob[p + 4 : p + 4 + hl])
|
|
134
|
+
if hdr["version"] != VERSION or hdr["codec"] != "zstd":
|
|
135
|
+
raise ValueError(f"LP2 version {hdr['version']} codec {hdr['codec']}")
|
|
136
|
+
return hdr, p + 4 + hl
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def decode(blob: bytes) -> Decoded:
|
|
140
|
+
"""lobpack's `unpack2`, stopping before the list columns are rebuilt."""
|
|
141
|
+
hdr, p = read_header(blob)
|
|
142
|
+
codec = pa.Codec("zstd")
|
|
143
|
+
S = {}
|
|
144
|
+
for s in hdr["streams"]:
|
|
145
|
+
z = blob[p : p + s["z"]]
|
|
146
|
+
p += s["z"]
|
|
147
|
+
raw = codec.decompress(z, decompressed_size=s["raw"], asbytes=True) if s["raw"] else b""
|
|
148
|
+
S[s["name"]] = vdec(raw, s["count"])
|
|
149
|
+
if p != len(blob):
|
|
150
|
+
raise ValueError("trailing bytes after the last stream")
|
|
151
|
+
depth, names, m = hdr["depth"], hdr["fields"], hdr["rows_kept"]
|
|
152
|
+
run_tok = np.cumsum(S["run_tok"].astype(np.int64) + 1) - 1
|
|
153
|
+
run_len = S["run_len"].astype(np.int64) + 1
|
|
154
|
+
if int(run_len.sum()) != m:
|
|
155
|
+
raise ValueError("run lengths do not add up to rows_kept")
|
|
156
|
+
run_start = np.r_[0, np.cumsum(run_len)[:-1]].astype(np.int64)
|
|
157
|
+
new = np.zeros(m, bool)
|
|
158
|
+
new[run_start] = True
|
|
159
|
+
F = {"tok": np.repeat(run_tok, run_len)}
|
|
160
|
+
dts = np.zeros(m, np.int64)
|
|
161
|
+
dts[~new] = unzz(S["dts"])
|
|
162
|
+
dts[run_start] = unzz(S["run_ts0"]) + hdr["base_ts"]
|
|
163
|
+
F["ts_ms"] = seg_cumsum(dts, run_start, run_len)
|
|
164
|
+
F["ts_ex_ms"] = F["ts_ms"] - unzz(S["ex_lag"])
|
|
165
|
+
mask = S["mask"]
|
|
166
|
+
for bit, nm in enumerate(names):
|
|
167
|
+
ch = ((mask >> np.uint64(bit)) & np.uint64(1)).astype(bool)
|
|
168
|
+
d = np.zeros(m, np.int64)
|
|
169
|
+
d[ch] = unzz(S["f_" + nm])
|
|
170
|
+
F[nm] = seg_cumsum(d, run_start, run_len)
|
|
171
|
+
nf = len(names)
|
|
172
|
+
F["snap"] = ((mask >> np.uint64(nf)) & np.uint64(1)).astype(np.int64)
|
|
173
|
+
ti = np.flatnonzero(((mask >> np.uint64(nf + 1)) & np.uint64(1)).astype(bool))
|
|
174
|
+
if ti.size != hdr["rows_mirror"]:
|
|
175
|
+
raise ValueError(f"twin bits {ti.size} != rows_mirror {hdr['rows_mirror']}")
|
|
176
|
+
G = mirror_of(F, depth, ti)
|
|
177
|
+
for c in F:
|
|
178
|
+
F[c] = np.r_[F[c], G[c]]
|
|
179
|
+
o = np.lexsort((F["tok"], F["ts_ms"]))
|
|
180
|
+
for c in F:
|
|
181
|
+
F[c] = F[c][o]
|
|
182
|
+
if F["ts_ms"].size != hdr["rows"]:
|
|
183
|
+
raise ValueError("row count mismatch")
|
|
184
|
+
return Decoded(header=hdr, depth=depth, columns=F)
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
def to_table(decoded: Decoded) -> pa.Table:
|
|
188
|
+
"""lobpack's `from_arrays`: the books schema with list columns, as the recorder defines it."""
|
|
189
|
+
F, depth = decoded.columns, decoded.depth
|
|
190
|
+
cols = {}
|
|
191
|
+
for c, dt, cnt in (
|
|
192
|
+
("bid_px", pa.int32(), "nb"),
|
|
193
|
+
("bid_sz", pa.int64(), "nb"),
|
|
194
|
+
("ask_px", pa.int32(), "na"),
|
|
195
|
+
("ask_sz", pa.int64(), "na"),
|
|
196
|
+
):
|
|
197
|
+
cols[c] = pa.ListArray.from_arrays(
|
|
198
|
+
list_offsets(F[cnt]), level_values(F, c, depth, F[cnt]).astype(dt.to_pandas_dtype())
|
|
199
|
+
).cast(SCHEMA.field(c).type)
|
|
200
|
+
arrs = [
|
|
201
|
+
pa.array(F["ts_ms"], pa.int64()),
|
|
202
|
+
pa.array(F["ts_ex_ms"], pa.int64()),
|
|
203
|
+
pa.array(F["tok"].astype(np.int32), pa.int32()),
|
|
204
|
+
cols["bid_px"],
|
|
205
|
+
cols["bid_sz"],
|
|
206
|
+
cols["ask_px"],
|
|
207
|
+
cols["ask_sz"],
|
|
208
|
+
pa.array(F["snap"].astype(bool), pa.bool_()),
|
|
209
|
+
pa.array(F["shard"].astype(np.int32), pa.int32()),
|
|
210
|
+
pa.array(F["conn"].astype(np.int32), pa.int32()),
|
|
211
|
+
]
|
|
212
|
+
return pa.Table.from_arrays(arrs, schema=SCHEMA)
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
def list_offsets(lengths: np.ndarray) -> pa.Array:
|
|
216
|
+
"""List offsets for per-row level counts."""
|
|
217
|
+
return pa.array(np.r_[0, np.cumsum(lengths)].astype(np.int32), pa.int32())
|
|
218
|
+
|
|
219
|
+
|
|
220
|
+
def level_values(
|
|
221
|
+
F: dict[str, np.ndarray], side: str, depth: int, lengths: np.ndarray
|
|
222
|
+
) -> np.ndarray:
|
|
223
|
+
"""The present levels of one side, flattened row by row, best level first."""
|
|
224
|
+
if not depth:
|
|
225
|
+
return np.zeros(0, np.int64)
|
|
226
|
+
M = np.stack([F[f"{side}{k}"] for k in range(depth)], axis=1)
|
|
227
|
+
keep = np.arange(depth)[None, :] < lengths[:, None]
|
|
228
|
+
return M[keep]
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
# ---------- streaming: one window in bounded memory ----------
|
|
232
|
+
#
|
|
233
|
+
# Not in lobpack. A real hour is one window of about 35 million rows, and `decode` holds all of
|
|
234
|
+
# them at once (about 3 GB). The ingest runs on a box with far less, so it reads the window the
|
|
235
|
+
# way it is laid out: token by token. Every stream is decompressed incrementally and its varints
|
|
236
|
+
# are consumed through a cursor, a batch of whole token runs at a time. A batch never splits the
|
|
237
|
+
# two tokens of one market, so a mirrored row always lands in the same batch as its partner and a
|
|
238
|
+
# market's rows stay contiguous in the output.
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
class _Varints:
|
|
242
|
+
"""The varints of one zstd stream, read incrementally. `take(n)` returns the next n values."""
|
|
243
|
+
|
|
244
|
+
def __init__(self, name: str, frame: pa.Buffer, count: int, raw: int) -> None:
|
|
245
|
+
self._name = name
|
|
246
|
+
self._stream = pa.input_stream(frame, compression="zstd") if raw else None
|
|
247
|
+
self._buf = np.zeros(0, np.uint8)
|
|
248
|
+
self._left = count
|
|
249
|
+
self._raw_left = raw
|
|
250
|
+
|
|
251
|
+
def take(self, n: int) -> np.ndarray:
|
|
252
|
+
if n > self._left:
|
|
253
|
+
raise ValueError(f"stream {self._name}: {n} values asked, {self._left} left")
|
|
254
|
+
if n == 0:
|
|
255
|
+
return np.zeros(0, np.uint64)
|
|
256
|
+
while True:
|
|
257
|
+
ends = np.flatnonzero(self._buf < 0x80)
|
|
258
|
+
if ends.size >= n:
|
|
259
|
+
break
|
|
260
|
+
if self._stream is None:
|
|
261
|
+
raise ValueError(f"stream {self._name} ended early")
|
|
262
|
+
more = self._stream.read(max(1 << 20, 2 * (n - ends.size)))
|
|
263
|
+
if not more:
|
|
264
|
+
raise ValueError(f"stream {self._name} ended early")
|
|
265
|
+
self._raw_left -= len(more)
|
|
266
|
+
self._buf = np.concatenate([self._buf, np.frombuffer(more, np.uint8)])
|
|
267
|
+
cut = int(ends[n - 1]) + 1
|
|
268
|
+
values = vdec(self._buf[:cut].tobytes(), n)
|
|
269
|
+
self._buf = self._buf[cut:]
|
|
270
|
+
self._left -= n
|
|
271
|
+
return values
|
|
272
|
+
|
|
273
|
+
def finish(self) -> None:
|
|
274
|
+
"""Every value consumed and not a byte more - the streaming form of lobpack's checks."""
|
|
275
|
+
trailing = self._stream.read() if self._stream is not None else b""
|
|
276
|
+
self._raw_left -= len(trailing)
|
|
277
|
+
if self._left or self._buf.size or trailing or self._raw_left:
|
|
278
|
+
raise ValueError(f"stream {self._name} does not match its header")
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
def stream(buffer: pa.Buffer, max_rows: int) -> Iterator[Decoded]:
|
|
282
|
+
"""The rows of one LP2 object in batches of whole markets, each in (tok, ts_ms) order.
|
|
283
|
+
|
|
284
|
+
`max_rows` bounds the stored rows of a batch; mirrored rows come on top, at most as many
|
|
285
|
+
again. A single market larger than that is still one batch, because splitting it would
|
|
286
|
+
separate a row from its mirror. Within a market the order is lobpack's tie order.
|
|
287
|
+
"""
|
|
288
|
+
hdr, p = read_header(buffer.slice(0, min(buffer.size, 1 << 20)).to_pybytes())
|
|
289
|
+
cursors: dict[str, _Varints] = {}
|
|
290
|
+
for s in hdr["streams"]:
|
|
291
|
+
if p + s["z"] > buffer.size:
|
|
292
|
+
raise ValueError(f"stream {s['name']} runs past the end of the object")
|
|
293
|
+
cursors[s["name"]] = _Varints(s["name"], buffer.slice(p, s["z"]), s["count"], s["raw"])
|
|
294
|
+
p += s["z"]
|
|
295
|
+
if p != buffer.size:
|
|
296
|
+
raise ValueError("trailing bytes after the last stream")
|
|
297
|
+
depth, names = hdr["depth"], hdr["fields"]
|
|
298
|
+
runs = next(s["count"] for s in hdr["streams"] if s["name"] == "run_tok")
|
|
299
|
+
run_tok = np.cumsum(cursors["run_tok"].take(runs).astype(np.int64) + 1) - 1
|
|
300
|
+
run_len = cursors["run_len"].take(runs).astype(np.int64) + 1
|
|
301
|
+
run_ts0 = unzz(cursors["run_ts0"].take(runs)) + hdr["base_ts"]
|
|
302
|
+
if int(run_len.sum()) != hdr["rows_kept"]:
|
|
303
|
+
raise ValueError("run lengths do not add up to rows_kept")
|
|
304
|
+
|
|
305
|
+
# A batch may end only where the next run starts another market (tok // 2 changes).
|
|
306
|
+
market_start = np.r_[True, (run_tok[1:] // 2) != (run_tok[:-1] // 2)]
|
|
307
|
+
rows_seen = mirrors_seen = 0
|
|
308
|
+
a = 0
|
|
309
|
+
while a < runs:
|
|
310
|
+
b = a + 1
|
|
311
|
+
total = int(run_len[a])
|
|
312
|
+
while b < runs and (not market_start[b] or total + int(run_len[b]) <= max_rows):
|
|
313
|
+
total += int(run_len[b])
|
|
314
|
+
b += 1
|
|
315
|
+
lens = run_len[a:b]
|
|
316
|
+
m = int(lens.sum())
|
|
317
|
+
starts = np.r_[0, np.cumsum(lens)[:-1]].astype(np.int64)
|
|
318
|
+
new = np.zeros(m, bool)
|
|
319
|
+
new[starts] = True
|
|
320
|
+
F = {"tok": np.repeat(run_tok[a:b], lens)}
|
|
321
|
+
dts = np.zeros(m, np.int64)
|
|
322
|
+
dts[~new] = unzz(cursors["dts"].take(m - (b - a)))
|
|
323
|
+
dts[starts] = run_ts0[a:b]
|
|
324
|
+
F["ts_ms"] = seg_cumsum(dts, starts, lens)
|
|
325
|
+
F["ts_ex_ms"] = F["ts_ms"] - unzz(cursors["ex_lag"].take(m))
|
|
326
|
+
mask = cursors["mask"].take(m)
|
|
327
|
+
for bit, nm in enumerate(names):
|
|
328
|
+
ch = ((mask >> np.uint64(bit)) & np.uint64(1)).astype(bool)
|
|
329
|
+
d = np.zeros(m, np.int64)
|
|
330
|
+
d[ch] = unzz(cursors["f_" + nm].take(int(ch.sum())))
|
|
331
|
+
F[nm] = seg_cumsum(d, starts, lens)
|
|
332
|
+
nf = len(names)
|
|
333
|
+
F["snap"] = ((mask >> np.uint64(nf)) & np.uint64(1)).astype(np.int64)
|
|
334
|
+
ti = np.flatnonzero(((mask >> np.uint64(nf + 1)) & np.uint64(1)).astype(bool))
|
|
335
|
+
G = mirror_of(F, depth, ti)
|
|
336
|
+
for c in F:
|
|
337
|
+
F[c] = np.r_[F[c], G[c]]
|
|
338
|
+
o = np.lexsort((F["ts_ms"], F["tok"]))
|
|
339
|
+
for c in F:
|
|
340
|
+
F[c] = F[c][o]
|
|
341
|
+
rows_seen += F["ts_ms"].size
|
|
342
|
+
mirrors_seen += ti.size
|
|
343
|
+
yield Decoded(header=hdr, depth=depth, columns=F)
|
|
344
|
+
a = b
|
|
345
|
+
|
|
346
|
+
# Streams this reader decodes are held to their header exactly. A stream it does not know - an
|
|
347
|
+
# addition a newer writer makes, like the legacy archive's sub_ns - is skipped, the way
|
|
348
|
+
# lobpack's unpack2 skips it: an extension must not make an older reader refuse the file.
|
|
349
|
+
known = {"run_tok", "run_len", "run_ts0", "dts", "ex_lag", "mask"} | {f"f_{n}" for n in names}
|
|
350
|
+
for name, cursor in cursors.items():
|
|
351
|
+
if name in known:
|
|
352
|
+
cursor.finish()
|
|
353
|
+
if rows_seen != hdr["rows"] or mirrors_seen != hdr["rows_mirror"]:
|
|
354
|
+
raise ValueError(
|
|
355
|
+
f"rows {rows_seen} (mirror {mirrors_seen}) != header {hdr['rows']} "
|
|
356
|
+
f"({hdr['rows_mirror']})"
|
|
357
|
+
)
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
"""Open any downloaded ProbNode file as a pyarrow Table: lp2, parquet or csv.gz."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from collections.abc import Iterator
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
import pyarrow as pa
|
|
9
|
+
|
|
10
|
+
#: Stored rows per decode batch. About 0.5 GB of working memory on a real hour of books; a
|
|
11
|
+
#: whole-object decode of the same hour needs ~8 GB, which is why lp2 is always read in batches.
|
|
12
|
+
LP2_BATCH_ROWS = 200_000
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def iter_lp2(path: str | Path, batch_rows: int = LP2_BATCH_ROWS) -> Iterator[pa.Table]:
|
|
16
|
+
"""The books in one lp2 file as tables of whole markets, in bounded memory.
|
|
17
|
+
|
|
18
|
+
Columns: ts_ms (our receive clock) and ts_ex_ms (the venue's), tok (index into the file's
|
|
19
|
+
token list), the bid/ask price and size lists - prices x10000, sizes x1e6, best level first -
|
|
20
|
+
snap, shard and conn.
|
|
21
|
+
"""
|
|
22
|
+
from . import lp2
|
|
23
|
+
|
|
24
|
+
buffer = pa.memory_map(str(path)).read_buffer()
|
|
25
|
+
for decoded in lp2.stream(buffer, batch_rows):
|
|
26
|
+
yield lp2.to_table(decoded)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def read(path: str | Path) -> pa.Table:
|
|
30
|
+
"""A pyarrow Table for one file. `.to_pandas()` or `polars.from_arrow()` from there.
|
|
31
|
+
|
|
32
|
+
An hour of books is tens of millions of rows; for lp2 prefer `iter_lp2` unless the whole hour
|
|
33
|
+
fits in memory, and for parquet `pyarrow.parquet.ParquetFile(path).iter_batches()`.
|
|
34
|
+
"""
|
|
35
|
+
path = Path(path)
|
|
36
|
+
name = path.name
|
|
37
|
+
if name.endswith(".lp2"):
|
|
38
|
+
return pa.concat_tables(iter_lp2(path))
|
|
39
|
+
if name.endswith(".parquet"):
|
|
40
|
+
import pyarrow.parquet as pq
|
|
41
|
+
|
|
42
|
+
return pq.read_table(path)
|
|
43
|
+
if name.endswith(".csv.gz") or name.endswith(".csv"):
|
|
44
|
+
import pyarrow.csv as pacsv
|
|
45
|
+
|
|
46
|
+
return pacsv.read_csv(path)
|
|
47
|
+
raise ValueError(f"not a ProbNode data file (lp2, parquet, csv.gz): {path}")
|
|
@@ -0,0 +1,204 @@
|
|
|
1
|
+
"""The SDK and CLI against a fake ProbNode: an httpx MockTransport that serves a listing, signs
|
|
2
|
+
paths and hands out the bytes behind each signed URL."""
|
|
3
|
+
|
|
4
|
+
from __future__ import annotations
|
|
5
|
+
|
|
6
|
+
import hashlib
|
|
7
|
+
import json
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
import httpx
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
import probnode
|
|
14
|
+
from probnode import cli
|
|
15
|
+
from probnode.client import Client, ProbNodeError
|
|
16
|
+
|
|
17
|
+
BLOBS = {
|
|
18
|
+
"x/parquet/books/event_date=2026-11-08/hour=12.parquet": b"twelve" * 1000,
|
|
19
|
+
"x/parquet/books/event_date=2026-11-08/hour=13.parquet": b"thirteen" * 1000,
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class Fake:
|
|
24
|
+
def __init__(self) -> None:
|
|
25
|
+
self.blobs = dict(BLOBS)
|
|
26
|
+
self.fetches: list[str] = []
|
|
27
|
+
self.signed: list[str] = []
|
|
28
|
+
|
|
29
|
+
def listing(self) -> list[dict]:
|
|
30
|
+
return [
|
|
31
|
+
{
|
|
32
|
+
"dataset": "books",
|
|
33
|
+
"format": "parquet",
|
|
34
|
+
"event_date": "2026-11-08",
|
|
35
|
+
"hour": int(p.split("hour=")[1][:2]),
|
|
36
|
+
"src": "s6",
|
|
37
|
+
"object_path": p,
|
|
38
|
+
"rows": 10,
|
|
39
|
+
"bytes": len(b),
|
|
40
|
+
"sha256": hashlib.sha256(b).hexdigest(),
|
|
41
|
+
"granularity": "hour",
|
|
42
|
+
"market": None,
|
|
43
|
+
"future_field": 1,
|
|
44
|
+
}
|
|
45
|
+
for p, b in self.blobs.items()
|
|
46
|
+
]
|
|
47
|
+
|
|
48
|
+
def __call__(self, request: httpx.Request) -> httpx.Response:
|
|
49
|
+
url = request.url
|
|
50
|
+
if url.host == "lake.test":
|
|
51
|
+
path = url.path.removeprefix("/")
|
|
52
|
+
self.fetches.append(path)
|
|
53
|
+
return httpx.Response(200, content=self.served(path))
|
|
54
|
+
if request.headers.get("Authorization") != "Bearer k":
|
|
55
|
+
return httpx.Response(401, json={"detail": "missing or invalid API key"})
|
|
56
|
+
if url.path == "/v1/data/files":
|
|
57
|
+
if url.params.get("dataset") == "nope":
|
|
58
|
+
return httpx.Response(400, json={"detail": "unknown dataset 'nope'"})
|
|
59
|
+
files = self.listing()
|
|
60
|
+
return httpx.Response(
|
|
61
|
+
200,
|
|
62
|
+
json={
|
|
63
|
+
"files": files,
|
|
64
|
+
"total_files": len(files),
|
|
65
|
+
"total_bytes": sum(f["bytes"] for f in files),
|
|
66
|
+
},
|
|
67
|
+
)
|
|
68
|
+
if url.path == "/v1/data/urls":
|
|
69
|
+
paths = json.loads(request.content)["paths"]
|
|
70
|
+
self.signed += paths
|
|
71
|
+
return httpx.Response(
|
|
72
|
+
200,
|
|
73
|
+
json={
|
|
74
|
+
"urls": [
|
|
75
|
+
{"object_path": p, "url": f"https://lake.test/{p}", "expires_at": "x"}
|
|
76
|
+
for p in paths
|
|
77
|
+
],
|
|
78
|
+
"licence": "research",
|
|
79
|
+
"licence_detail": "",
|
|
80
|
+
},
|
|
81
|
+
)
|
|
82
|
+
return httpx.Response(404, json={"detail": "not found"})
|
|
83
|
+
|
|
84
|
+
def served(self, path: str) -> bytes:
|
|
85
|
+
return self.blobs[path]
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
@pytest.fixture
|
|
89
|
+
def fake(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Fake:
|
|
90
|
+
f = Fake()
|
|
91
|
+
monkeypatch.setenv("PROBNODE_CONFIG", str(tmp_path / "cfg.json"))
|
|
92
|
+
real_init = Client.__init__
|
|
93
|
+
|
|
94
|
+
def init(self, api_key=None, base_url=None, **kw):
|
|
95
|
+
real_init(
|
|
96
|
+
self,
|
|
97
|
+
api_key=api_key or "k",
|
|
98
|
+
base_url=base_url or "https://api.test",
|
|
99
|
+
transport=httpx.MockTransport(f),
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
monkeypatch.setattr(Client, "__init__", init)
|
|
103
|
+
return f
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
ARGS = ["books", "parquet", "2026-11-08T12", "2026-11-08T14"]
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def test_files_keeps_unknown_fields_and_parses(fake: Fake) -> None:
|
|
110
|
+
files = Client().files("books", "parquet", "2026-11-08T12:00:00Z", "2026-11-08T14:00:00Z")
|
|
111
|
+
assert [f.hour for f in files] == [12, 13]
|
|
112
|
+
assert files[0].extra == {"future_field": 1}
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def test_api_refusal_is_the_api_sentence(fake: Fake) -> None:
|
|
116
|
+
with pytest.raises(ProbNodeError) as err:
|
|
117
|
+
Client().files("nope", "parquet", "a", "b")
|
|
118
|
+
assert err.value.status == 400 and "unknown dataset" in err.value.detail
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def test_download_verifies_sha256(fake: Fake, tmp_path: Path) -> None:
|
|
122
|
+
pn = Client()
|
|
123
|
+
f = pn.files(*ARGS[:2], "s", "e")[0]
|
|
124
|
+
fake.served = lambda path: b"tampered" # type: ignore[method-assign]
|
|
125
|
+
with pytest.raises(ProbNodeError, match="sha256 mismatch"):
|
|
126
|
+
pn.download(f, tmp_path)
|
|
127
|
+
assert not (tmp_path / f.object_path).exists()
|
|
128
|
+
assert not list(tmp_path.rglob("*.part"))
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
def test_sync_twice_moves_nothing_then_only_the_change(fake: Fake, tmp_path: Path) -> None:
|
|
132
|
+
out = tmp_path / "data"
|
|
133
|
+
assert cli.main(["sync", *ARGS, "-o", str(out)]) == 0
|
|
134
|
+
assert len(fake.fetches) == 2
|
|
135
|
+
assert cli.main(["sync", *ARGS, "-o", str(out)]) == 0
|
|
136
|
+
assert len(fake.fetches) == 2, "second sync fetched again"
|
|
137
|
+
|
|
138
|
+
changed = "x/parquet/books/event_date=2026-11-08/hour=13.parquet"
|
|
139
|
+
fake.blobs[changed] = b"corrected" * 1000
|
|
140
|
+
assert cli.main(["sync", *ARGS, "-o", str(out)]) == 0
|
|
141
|
+
assert fake.fetches[2:] == [changed]
|
|
142
|
+
assert (out / changed).read_bytes() == fake.blobs[changed]
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def test_doctor_names_a_corrupt_and_a_missing_file(
|
|
146
|
+
fake: Fake, tmp_path: Path, capsys: pytest.CaptureFixture[str]
|
|
147
|
+
) -> None:
|
|
148
|
+
out = tmp_path / "data"
|
|
149
|
+
cli.main(["sync", *ARGS, "-o", str(out)])
|
|
150
|
+
assert cli.main(["doctor", str(out)]) == 0
|
|
151
|
+
paths = sorted(BLOBS)
|
|
152
|
+
(out / paths[0]).write_bytes(b"bitrot")
|
|
153
|
+
(out / paths[1]).unlink()
|
|
154
|
+
capsys.readouterr()
|
|
155
|
+
assert cli.main(["doctor", str(out)]) == 1
|
|
156
|
+
text = capsys.readouterr().out
|
|
157
|
+
assert f"CORRUPT {paths[0]}" in text and f"MISSING {paths[1]}" in text
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
def test_cli_refusal_exits_2(fake: Fake, capsys: pytest.CaptureFixture[str]) -> None:
|
|
161
|
+
assert cli.main(["files", "nope", "parquet", "2026-11-08T12", "2026-11-08T13"]) == 2
|
|
162
|
+
assert "unknown dataset" in capsys.readouterr().err
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def test_iso_hour_accepts_the_short_forms() -> None:
|
|
166
|
+
assert cli.iso_hour("2026-11-08T12") == "2026-11-08T12:00:00Z"
|
|
167
|
+
assert cli.iso_hour("2026-11-08 07:00") == "2026-11-08T07:00:00Z"
|
|
168
|
+
assert cli.iso_hour("2026-11-08T12:00:00Z") == "2026-11-08T12:00:00Z"
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
def test_read_parquet_and_csv(tmp_path: Path) -> None:
|
|
172
|
+
import pyarrow as pa
|
|
173
|
+
import pyarrow.csv as pacsv
|
|
174
|
+
import pyarrow.parquet as pq
|
|
175
|
+
|
|
176
|
+
t = pa.table({"a": [1, 2]})
|
|
177
|
+
pq.write_table(t, tmp_path / "x.parquet")
|
|
178
|
+
pacsv.write_csv(t, tmp_path / "x.csv")
|
|
179
|
+
assert probnode.read(tmp_path / "x.parquet").to_pydict() == {"a": [1, 2]}
|
|
180
|
+
assert probnode.read(tmp_path / "x.csv").to_pydict() == {"a": [1, 2]}
|
|
181
|
+
with pytest.raises(ValueError):
|
|
182
|
+
probnode.read(tmp_path / "x.txt")
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
def test_lp2_reader_is_the_api_one() -> None:
|
|
186
|
+
"""The SDK ships the same lp2 decoder the ingest uses; if they drift, a customer decodes
|
|
187
|
+
differently from the exports we made. Byte-identical, checked where both are present."""
|
|
188
|
+
api = Path(__file__).resolve().parents[3] / "apps/api/src/pmd_api/lp2.py"
|
|
189
|
+
if not api.exists():
|
|
190
|
+
pytest.skip("not in the monorepo")
|
|
191
|
+
ours = Path(probnode.__file__).with_name("lp2.py")
|
|
192
|
+
assert ours.read_bytes() == api.read_bytes()
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def test_sync_refetches_a_file_changed_on_disk(fake: Fake, tmp_path: Path) -> None:
|
|
196
|
+
"""The state file says current, but the bytes on disk were changed: sync must notice and fetch
|
|
197
|
+
again, not trust its own bookkeeping."""
|
|
198
|
+
out = tmp_path / "data"
|
|
199
|
+
cli.main(["sync", *ARGS, "-o", str(out)])
|
|
200
|
+
victim = sorted(BLOBS)[0]
|
|
201
|
+
(out / victim).write_bytes(b"corrupt")
|
|
202
|
+
assert cli.main(["sync", *ARGS, "-o", str(out)]) == 0
|
|
203
|
+
assert fake.fetches[2:] == [victim]
|
|
204
|
+
assert cli.main(["doctor", str(out)]) == 0
|