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.
@@ -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