global: snapshot

This commit is contained in:
nym21
2026-02-24 12:21:20 +01:00
parent 3b7aa8242a
commit cefc8cfd42
24 changed files with 3561 additions and 1073 deletions

View File

@@ -99,16 +99,16 @@ class BrkClientBase:
"""Make a GET request and return text."""
return self.get(path).decode()
def close(self):
def close(self) -> None:
"""Close the HTTP client."""
if self._conn:
self._conn.close()
self._conn = None
def __enter__(self):
def __enter__(self) -> BrkClientBase:
return self
def __exit__(self, exc_type, exc_val, exc_tb):
def __exit__(self, exc_type: Optional[type], exc_val: Optional[BaseException], exc_tb: Optional[Any]) -> None:
self.close()
@@ -221,7 +221,7 @@ def _date_to_index(index: str, d: Union[date, datetime]) -> int:
@dataclass
class MetricData(Generic[T]):
"""Metric data with range information."""
"""Metric data with range information. Always int-indexed."""
version: int
index: Index
total: int
@@ -235,63 +235,96 @@ class MetricData(Generic[T]):
"""Whether this metric uses a date-based index."""
return self.index in _DATE_INDEXES
def dates(self) -> list:
"""Get dates for the index range. Date-based indexes only, throws otherwise."""
return [_index_to_date(self.index, i) for i in range(self.start, self.end)]
def indexes(self) -> List[int]:
"""Get raw index numbers."""
return list(range(self.start, self.end))
def keys(self) -> list:
"""Get keys: dates for date-based indexes, index numbers otherwise."""
return self.dates() if self.is_date_based else self.indexes()
def keys(self) -> List[int]:
"""Get keys as index numbers."""
return self.indexes()
def items(self) -> list:
"""Get (key, value) pairs: keys are dates for date-based, numbers otherwise."""
return list(zip(self.keys(), self.data))
def items(self) -> List[Tuple[int, T]]:
"""Get (index, value) pairs."""
return list(zip(self.indexes(), self.data))
def to_dict(self) -> dict:
"""Return {{key: value}} dict: keys are dates for date-based, numbers otherwise."""
return dict(zip(self.keys(), self.data))
def to_dict(self) -> Dict[int, T]:
"""Return {{index: value}} dict."""
return dict(zip(self.indexes(), self.data))
def __iter__(self):
"""Iterate over (key, value) pairs. Keys are dates for date-based, numbers otherwise."""
return iter(zip(self.keys(), self.data))
def __iter__(self) -> Iterator[Tuple[int, T]]:
"""Iterate over (index, value) pairs."""
return iter(zip(self.indexes(), self.data))
def __len__(self) -> int:
return len(self.data)
def to_polars(self) -> pl.DataFrame:
"""Convert to Polars DataFrame with 'index' and 'value' columns."""
try:
import polars as pl # type: ignore[import-not-found]
except ImportError:
raise ImportError("polars is required: pip install polars")
return pl.DataFrame({{"index": self.indexes(), "value": self.data}})
def to_pandas(self) -> pd.DataFrame:
"""Convert to Pandas DataFrame with 'index' and 'value' columns."""
try:
import pandas as pd # type: ignore[import-not-found]
except ImportError:
raise ImportError("pandas is required: pip install pandas")
return pd.DataFrame({{"index": self.indexes(), "value": self.data}})
@dataclass
class DateMetricData(MetricData[T]):
"""Metric data with date-based index. Extends MetricData with date methods."""
def dates(self) -> List[Union[date, datetime]]:
"""Get dates for the index range. Returns datetime for sub-daily indexes, date for daily+."""
return [_index_to_date(self.index, i) for i in range(self.start, self.end)]
def date_items(self) -> List[Tuple[Union[date, datetime], T]]:
"""Get (date, value) pairs."""
return list(zip(self.dates(), self.data))
def to_date_dict(self) -> Dict[Union[date, datetime], T]:
"""Return {{date: value}} dict."""
return dict(zip(self.dates(), self.data))
def to_polars(self, with_dates: bool = True) -> pl.DataFrame:
"""Convert to Polars DataFrame. Requires polars to be installed.
"""Convert to Polars DataFrame.
Returns a DataFrame with columns:
- 'date' and 'value' if with_dates=True and index is date-based
- 'date' and 'value' if with_dates=True (default)
- 'index' and 'value' otherwise
"""
try:
import polars as pl # type: ignore[import-not-found]
except ImportError:
raise ImportError("polars is required: pip install polars")
if with_dates and self.is_date_based:
if with_dates:
return pl.DataFrame({{"date": self.dates(), "value": self.data}})
return pl.DataFrame({{"index": self.indexes(), "value": self.data}})
def to_pandas(self, with_dates: bool = True) -> pd.DataFrame:
"""Convert to Pandas DataFrame. Requires pandas to be installed.
"""Convert to Pandas DataFrame.
Returns a DataFrame with columns:
- 'date' and 'value' if with_dates=True and index is date-based
- 'date' and 'value' if with_dates=True (default)
- 'index' and 'value' otherwise
"""
try:
import pandas as pd # type: ignore[import-not-found]
except ImportError:
raise ImportError("pandas is required: pip install pandas")
if with_dates and self.is_date_based:
if with_dates:
return pd.DataFrame({{"date": self.dates(), "value": self.data}})
return pd.DataFrame({{"index": self.indexes(), "value": self.data}})
# Type alias for non-generic usage
# Type aliases for non-generic usage
AnyMetricData = MetricData[Any]
AnyDateMetricData = DateMetricData[Any]
class _EndpointConfig:
@@ -325,9 +358,15 @@ class _EndpointConfig:
p = self.path()
return f"{{p}}?{{query}}" if query else p
def get_metric(self) -> MetricData:
def _new(self, start: Optional[int] = None, end: Optional[int] = None) -> _EndpointConfig:
return _EndpointConfig(self.client, self.name, self.index, start, end)
def get_metric(self) -> MetricData[Any]:
return MetricData(**self.client.get_json(self._build_path()))
def get_date_metric(self) -> DateMetricData[Any]:
return DateMetricData(**self.client.get_json(self._build_path()))
def get_csv(self) -> str:
return self.client.get_text(self._build_path(format='csv'))
@@ -371,10 +410,7 @@ class SkippedBuilder(Generic[T]):
def take(self, n: int) -> RangeBuilder[T]:
"""Take n items after the skipped position."""
start = self._config.start or 0
return RangeBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
start, start + n
))
return RangeBuilder(self._config._new(start, start + n))
def fetch(self) -> MetricData[T]:
"""Fetch from skipped position to end."""
@@ -385,29 +421,35 @@ class SkippedBuilder(Generic[T]):
return self._config.get_csv()
class MetricEndpointBuilder(Generic[T]):
"""Builder for metric endpoint queries.
class DateRangeBuilder(RangeBuilder[T]):
"""Range builder that returns DateMetricData."""
def fetch(self) -> DateMetricData[T]:
return self._config.get_date_metric()
Use method chaining to specify the data range, then call fetch() or fetch_csv() to execute.
class DateSingleItemBuilder(SingleItemBuilder[T]):
"""Single item builder that returns DateMetricData."""
def fetch(self) -> DateMetricData[T]:
return self._config.get_date_metric()
class DateSkippedBuilder(SkippedBuilder[T]):
"""Skipped builder that returns DateMetricData."""
def take(self, n: int) -> DateRangeBuilder[T]:
start = self._config.start or 0
return DateRangeBuilder(self._config._new(start, start + n))
def fetch(self) -> DateMetricData[T]:
return self._config.get_date_metric()
class MetricEndpointBuilder(Generic[T]):
"""Builder for metric endpoint queries with int-based indexing.
Examples:
# Fetch all data
data = endpoint.fetch()
# Single item access
data = endpoint[5].fetch()
# Slice syntax (Python-native)
data = endpoint[:10].fetch() # First 10
data = endpoint[-5:].fetch() # Last 5
data = endpoint[100:110].fetch() # Range
# Convenience methods (pandas-style)
data = endpoint.head().fetch() # First 10 (default)
data = endpoint.head(20).fetch() # First 20
data = endpoint.tail(5).fetch() # Last 5
# Iterator-style chaining
data = endpoint[:10].fetch()
data = endpoint.head(20).fetch()
data = endpoint.skip(100).take(10).fetch()
"""
@@ -419,66 +461,30 @@ class MetricEndpointBuilder(Generic[T]):
@overload
def __getitem__(self, key: slice) -> RangeBuilder[T]: ...
def __getitem__(self, key: Union[int, slice, date, datetime]) -> Union[SingleItemBuilder[T], RangeBuilder[T]]:
"""Access single item or slice. Accepts dates for date-based indexes.
Examples:
endpoint[5] # Single item at index 5
endpoint[:10] # First 10
endpoint[-5:] # Last 5
endpoint[100:110] # Range 100-109
endpoint[date(2020, 1, 1):date(2023, 1, 1)] # Date range
endpoint[date(2020, 1, 1):] # Since date
"""
if isinstance(key, (date, datetime)):
idx = _date_to_index(self._config.index, key)
return SingleItemBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
idx, idx + 1
))
def __getitem__(self, key: Union[int, slice]) -> Union[SingleItemBuilder[T], RangeBuilder[T]]:
"""Access single item or slice by integer index."""
if isinstance(key, int):
return SingleItemBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
key, key + 1
))
start, stop = key.start, key.stop
if isinstance(start, (date, datetime)):
start = _date_to_index(self._config.index, start)
if isinstance(stop, (date, datetime)):
stop = _date_to_index(self._config.index, stop)
return RangeBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
start, stop
))
return SingleItemBuilder(self._config._new(key, key + 1))
return RangeBuilder(self._config._new(key.start, key.stop))
def head(self, n: int = 10) -> RangeBuilder[T]:
"""Get the first n items (pandas-style)."""
return RangeBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
None, n
))
"""Get the first n items."""
return RangeBuilder(self._config._new(end=n))
def tail(self, n: int = 10) -> RangeBuilder[T]:
"""Get the last n items (pandas-style)."""
start, end = (None, 0) if n == 0 else (-n, None)
return RangeBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
start, end
))
"""Get the last n items."""
return RangeBuilder(self._config._new(end=0) if n == 0 else self._config._new(start=-n))
def skip(self, n: int) -> SkippedBuilder[T]:
"""Skip the first n items. Chain with take() to get a range."""
return SkippedBuilder(_EndpointConfig(
self._config.client, self._config.name, self._config.index,
n, None
))
"""Skip the first n items."""
return SkippedBuilder(self._config._new(start=n))
def fetch(self) -> MetricData[T]:
"""Fetch all data as parsed JSON."""
"""Fetch all data."""
return self._config.get_metric()
def fetch_csv(self) -> str:
"""Fetch all data as CSV string."""
"""Fetch all data as CSV."""
return self._config.get_csv()
def path(self) -> str:
@@ -486,8 +492,72 @@ class MetricEndpointBuilder(Generic[T]):
return self._config.path()
# Type alias for non-generic usage
class DateMetricEndpointBuilder(Generic[T]):
"""Builder for metric endpoint queries with date-based indexing.
Accepts dates in __getitem__ and returns DateMetricData from fetch().
Examples:
data = endpoint.fetch()
data = endpoint[date(2020, 1, 1)].fetch()
data = endpoint[date(2020, 1, 1):date(2023, 1, 1)].fetch()
data = endpoint[:10].fetch()
"""
def __init__(self, client: BrkClientBase, name: str, index: Index):
self._config = _EndpointConfig(client, name, index)
@overload
def __getitem__(self, key: int) -> DateSingleItemBuilder[T]: ...
@overload
def __getitem__(self, key: datetime) -> DateSingleItemBuilder[T]: ...
@overload
def __getitem__(self, key: date) -> DateSingleItemBuilder[T]: ...
@overload
def __getitem__(self, key: slice) -> DateRangeBuilder[T]: ...
def __getitem__(self, key: Union[int, slice, date, datetime]) -> Union[DateSingleItemBuilder[T], DateRangeBuilder[T]]:
"""Access single item or slice. Accepts int, date, or datetime."""
if isinstance(key, (date, datetime)):
idx = _date_to_index(self._config.index, key)
return DateSingleItemBuilder(self._config._new(idx, idx + 1))
if isinstance(key, int):
return DateSingleItemBuilder(self._config._new(key, key + 1))
start, stop = key.start, key.stop
if isinstance(start, (date, datetime)):
start = _date_to_index(self._config.index, start)
if isinstance(stop, (date, datetime)):
stop = _date_to_index(self._config.index, stop)
return DateRangeBuilder(self._config._new(start, stop))
def head(self, n: int = 10) -> DateRangeBuilder[T]:
"""Get the first n items."""
return DateRangeBuilder(self._config._new(end=n))
def tail(self, n: int = 10) -> DateRangeBuilder[T]:
"""Get the last n items."""
return DateRangeBuilder(self._config._new(end=0) if n == 0 else self._config._new(start=-n))
def skip(self, n: int) -> DateSkippedBuilder[T]:
"""Skip the first n items."""
return DateSkippedBuilder(self._config._new(start=n))
def fetch(self) -> DateMetricData[T]:
"""Fetch all data."""
return self._config.get_date_metric()
def fetch_csv(self) -> str:
"""Fetch all data as CSV."""
return self._config.get_csv()
def path(self) -> str:
"""Get the base endpoint path."""
return self._config.path()
# Type aliases for non-generic usage
AnyMetricEndpointBuilder = MetricEndpointBuilder[Any]
AnyDateMetricEndpointBuilder = DateMetricEndpointBuilder[Any]
class MetricPattern(Protocol[T]):
@@ -535,11 +605,14 @@ pub fn generate_index_accessors(output: &mut String, patterns: &[IndexSetPattern
}
writeln!(output).unwrap();
// Generate helper function
// Generate helper functions
writeln!(
output,
r#"def _ep(c: BrkClientBase, n: str, i: Index) -> MetricEndpointBuilder[Any]:
return MetricEndpointBuilder(c, n, i)
def _dep(c: BrkClientBase, n: str, i: Index) -> DateMetricEndpointBuilder[Any]:
return DateMetricEndpointBuilder(c, n, i)
"#
)
.unwrap();
@@ -560,10 +633,15 @@ pub fn generate_index_accessors(output: &mut String, patterns: &[IndexSetPattern
for index in &pattern.indexes {
let method_name = index_to_field_name(index);
let index_name = index.name();
let (builder_type, helper) = if index.is_date_based() {
("DateMetricEndpointBuilder", "_dep")
} else {
("MetricEndpointBuilder", "_ep")
};
writeln!(
output,
" def {}(self) -> MetricEndpointBuilder[T]: return _ep(self._c, self._n, '{}')",
method_name, index_name
" def {}(self) -> {}[T]: return {}(self._c, self._n, '{}')",
method_name, builder_type, helper, index_name
)
.unwrap();
}