Skip to content

Commit

Permalink
Improve typing of read_json file and create common type alias
Browse files Browse the repository at this point in the history
Path type was never supported so also remove false type
  • Loading branch information
Zaczero committed Jan 27, 2025
1 parent f5f134b commit f879731
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions src/pypgstac/src/pypgstac/load.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,12 @@
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import (
Any,
BinaryIO,
Dict,
Generator,
Iterable,
Iterator,
Optional,
TextIO,
Tuple,
Expand Down Expand Up @@ -121,7 +119,10 @@ def open_std(
pass


def read_json(file: Union[Path, str, Iterator[Any]] = "stdin") -> Iterable:
_ReadJsonFileType = Union[str, Iterable[Union[Dict, bytes, bytearray, memoryview, str]]]


def read_json(file: _ReadJsonFileType = "stdin") -> Generator[Any, None, None]:
"""Load data from an ndjson or json file."""
if file is None:
file = "stdin"
Expand Down Expand Up @@ -205,7 +206,7 @@ def collection_json(self, collection_id: str) -> Tuple[Dict[str, Any], int, str]

def load_collections(
self,
file: Union[Path, str, Iterator[Any]] = "stdin",
file: _ReadJsonFileType = "stdin",
insert_mode: Optional[Methods] = Methods.insert,
) -> None:
"""Load a collections json or ndjson file."""
Expand Down Expand Up @@ -556,7 +557,10 @@ def _partition_update(self, item: Dict[str, Any]) -> str:

return partition_name

def read_dehydrated(self, file: Union[Path, str] = "stdin") -> Generator:
def read_dehydrated(
self,
file: str = "stdin",
) -> Generator[Dict[str, Any], None, None]:
if file is None:
file = "stdin"
if isinstance(file, str):
Expand Down Expand Up @@ -595,16 +599,16 @@ def read_dehydrated(self, file: Union[Path, str] = "stdin") -> Generator:

def read_hydrated(
self,
file: Union[Path, str, Iterator[Any]] = "stdin",
) -> Generator:
file: _ReadJsonFileType = "stdin",
) -> Generator[Dict[str, Any], None, None]:
for line in read_json(file):
item = self.format_item(line)
item["partition"] = self._partition_update(item)
yield item

def load_items(
self,
file: Union[Path, str, Iterator[Any]] = "stdin",
file: _ReadJsonFileType = "stdin",
insert_mode: Optional[Methods] = Methods.insert,
dehydrated: Optional[bool] = False,
chunksize: Optional[int] = 10000,
Expand All @@ -630,7 +634,7 @@ def load_items(

logger.debug(f"Adding data to database took {time.perf_counter() - t} seconds.")

def format_item(self, _item: Union[Path, str, Dict[str, Any]]) -> Dict[str, Any]:
def format_item(self, _item: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
"""Format an item to insert into a record."""
out: Dict[str, Any] = {}
item: Dict[str, Any]
Expand Down

0 comments on commit f879731

Please sign in to comment.