Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: add types to most of the uproot.source module #996

Merged
merged 10 commits into from
Oct 18, 2023
Merged
42 changes: 26 additions & 16 deletions src/uproot/source/chunk.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ class Source:
the file.
"""

def chunk(self, start, stop) -> Chunk:
def chunk(self, start: int, stop: int) -> Chunk:
"""
Args:
start (int): Seek position of the first byte to include.
Expand All @@ -60,7 +60,9 @@ def chunk(self, start, stop) -> Chunk:
:doc:`uproot.source.chunk.Chunk`.
"""

def chunks(self, ranges, notifications: queue.Queue) -> list[Chunk]:
def chunks(
self, ranges: list[(int, int)], notifications: queue.Queue
) -> list[Chunk]:
"""
Args:
ranges (list of (int, int) 2-tuples): Intervals to fetch
Expand Down Expand Up @@ -154,7 +156,7 @@ def __repr__(self):
type(self).__name__, path, self.num_workers, id(self)
)

def chunk(self, start, stop) -> Chunk:
def chunk(self, start: int, stop: int) -> Chunk:
self._num_requests += 1
self._num_requested_chunks += 1
self._num_requested_bytes += stop - start
Expand All @@ -164,7 +166,9 @@ def chunk(self, start, stop) -> Chunk:
self._executor.submit(future)
return chunk

def chunks(self, ranges, notifications: queue.Queue) -> list[Chunk]:
def chunks(
self, ranges: list[(int, int)], notifications: queue.Queue
) -> list[Chunk]:
self._num_requests += 1
self._num_requested_chunks += len(ranges)
self._num_requested_bytes += sum(stop - start for start, stop in ranges)
Expand All @@ -186,7 +190,7 @@ def executor(self):
return self._executor

@property
def num_workers(self):
def num_workers(self) -> int:
"""
The number of :doc:`uproot.source.futures.ResourceWorker` threads in
the :doc:`uproot.source.futures.ResourceThreadPoolExecutor`.
Expand Down Expand Up @@ -240,7 +244,7 @@ class Chunk:
_dtype = numpy.dtype(numpy.uint8)

@classmethod
def wrap(cls, source, data, start=0):
def wrap(cls, source: Source, data: numpy.ndarray, start: int = 0):
lobis marked this conversation as resolved.
Show resolved Hide resolved
"""
Args:
source (:doc:`uproot.source.chunk.Source`): Source to attach to
Expand All @@ -255,7 +259,9 @@ def wrap(cls, source, data, start=0):
future = uproot.source.futures.TrivialFuture(data)
return Chunk(source, start, start + len(data), future)

def __init__(self, source, start, stop, future, is_memmap=False):
def __init__(
self, source: Source, start: int, stop: int, future, is_memmap: bool = False
):
self._source = source
self._start = start
self._stop = stop
Expand All @@ -267,21 +273,21 @@ def __repr__(self):
return f"<Chunk {self._start}-{self._stop}>"

@property
def source(self):
def source(self) -> Source:
"""
Source from which this Chunk is derived.
"""
return self._source

@property
def start(self):
def start(self) -> int:
"""
Seek position of the first byte to include.
"""
return self._start

@property
def stop(self):
def stop(self) -> int:
"""
Seek position of the first byte to exclude (one greater than the last
byte to include).
Expand All @@ -297,7 +303,7 @@ def future(self):
return self._future

@property
def is_memmap(self):
def is_memmap(self) -> bool:
"""
If True, the `raw_data` is or will be a view into a memmap file, which
must be handled carefully. Accessing that data after the file is closed
Expand Down Expand Up @@ -326,15 +332,15 @@ def detach_memmap(self):
else:
return self

def __contains__(self, range):
def __contains__(self, range: tuple[int, int]):
start, stop = range
if isinstance(start, uproot.source.cursor.Cursor):
start = start.index
if isinstance(stop, uproot.source.cursor.Cursor):
stop = stop.index
return self._start <= start and stop <= self._stop

def wait(self, insist=True):
def wait(self, insist: bool = True):
"""
Args:
insist (bool or int): If True, raise an OSError if ``raw_data`` does
Expand Down Expand Up @@ -368,7 +374,7 @@ def wait(self, insist=True):
self._future = None

@property
def raw_data(self):
def raw_data(self) -> numpy.ndarray:
lobis marked this conversation as resolved.
Show resolved Hide resolved
"""
Data from the Source as a ``numpy.ndarray`` of ``numpy.uint8``.

Expand All @@ -379,7 +385,9 @@ def raw_data(self):
self.wait()
return self._raw_data

def get(self, start, stop, cursor, context):
def get(
self, start: int, stop: int, cursor: uproot.source.cursor.Cursor, context: dict
) -> numpy.ndarray:
"""
Args:
start (int): Seek position of the first byte to include.
Expand Down Expand Up @@ -417,7 +425,9 @@ def get(self, start, stop, cursor, context):
self._source.file_path,
)

def remainder(self, start, cursor, context):
def remainder(
self, start: int, cursor: uproot.source.cursor.Cursor, context: dict
) -> numpy.ndarray:
"""
Args:
start (int): Seek position of the first byte to include.
Expand Down
Loading