Skip to content

Commit

Permalink
a few type improvements, mostly flight related (#90)
Browse files Browse the repository at this point in the history
* FlightError.extra_info -> bytes

* annotate FlightStreamReader.cancel return

* BasicAuth serialize/deserialize

* RecordBatchFileReader.schema

* actually str | bytes
  • Loading branch information
dylanscott authored Sep 15, 2024
1 parent c0d156e commit 510e3bb
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
2 changes: 2 additions & 0 deletions pyarrow-stubs/__lib_pxi/ipc.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,8 @@ class _RecordBatchFileReader(_Weakrefable):
def __enter__(self) -> Self: ...
def __exit__(self, exc_type, exc_val, exc_tb): ...
@property
def schema(self) -> Schema: ...
@property
def stats(self) -> ReadStats: ...

def get_tensor_size(tensor: Tensor) -> int: ...
Expand Down
7 changes: 5 additions & 2 deletions pyarrow-stubs/_flight.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ class CertKeyPair(NamedTuple):
key: str

class FlightError(Exception):
extra_info: str
extra_info: bytes

class FlightInternalError(FlightError, ArrowException): ...
class FlightTimedOutError(FlightError, ArrowException): ...
Expand Down Expand Up @@ -83,6 +83,9 @@ class BasicAuth(_Weakrefable):
def username(self) -> bytes: ...
@property
def password(self) -> bytes: ...
def serialize(self) -> str: ...
@staticmethod
def deserialize(serialized: str | bytes) -> BasicAuth: ...

class DescriptorType(enum.Enum):
UNKNOWN = 0
Expand Down Expand Up @@ -196,7 +199,7 @@ class _MetadataRecordBatchReader(_Weakrefable, _ReadPandasMixin):
class MetadataRecordBatchReader(_MetadataRecordBatchReader): ...

class FlightStreamReader(MetadataRecordBatchReader):
def cancel(self): ...
def cancel(self) -> None: ...
def read_all(self) -> Table: ...

class MetadataRecordBatchWriter(_CRecordBatchWriter):
Expand Down

0 comments on commit 510e3bb

Please sign in to comment.