-
Notifications
You must be signed in to change notification settings - Fork 0
/
brin_parser.py
65 lines (53 loc) · 1.85 KB
/
brin_parser.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#%%
from __future__ import annotations
import csv
import re
from datetime import datetime
from typing import Iterable, Literal, Optional
from brin_lib import BlockRange, date_match
DT_TUPLE_REX = re.compile(
r"\{"
r"(?P<start>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})(?:\.\d+)?(?:[\+-]\d{2})?"
r" \.\. "
r"(?P<end>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})(?:\.\d+)?(?:[\+-]\d{2})?"
r"\}"
)
def parse_datetime_tuple(s: str) -> tuple[datetime, datetime]:
if (match := DT_TUPLE_REX.search(s)) is not None:
start, end = match.groups()
return (datetime.fromisoformat(start), datetime.fromisoformat(end))
else:
raise RuntimeError(f"failed to parse datetime tuple: {s}")
def parse_csv_rows(
csv_rows: Iterable[str],
start: Optional[datetime] = None,
end: Optional[datetime] = None,
match_type: Literal["overlap", "within"] = "overlap",
) -> list[BlockRange]:
reader = csv.reader(csv_rows)
brs = []
for i, row in enumerate(reader):
if i == 0:
assert row == ["blknum", "value"]
continue
blknum, value = row
br = BlockRange(int(blknum), *parse_datetime_tuple(value))
if date_match(br, match_type, start, end):
brs.append(br)
# note: not sorted because bug in export_brin_items.sh sorting different
# files. sorting makes viz look much nicer.
brs = sorted(brs, key=lambda br: br.blknum)
return brs
def parse_csv_file(
filepath: str,
start: Optional[datetime] = None,
end: Optional[datetime] = None,
match_type: Literal["overlap", "within"] = "overlap",
) -> list[BlockRange]:
with open(filepath) as f:
return parse_csv_rows(f, start, end, match_type)
#%%
if __name__ == "__main__":
blocks = parse_csv_file("brin_export_full.csv")
for i, block in enumerate(blocks[:10]):
print(f"{i:2d} {block}")