Add type hints to BlockPair

This commit is contained in:
2Shirt 2023-06-03 18:07:30 -07:00
parent fc2b90a2c0
commit 05de5c7294
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C

View file

@ -127,42 +127,42 @@ TIMEZONE = pytz.timezone(cfg.main.LINUX_TIME_ZONE)
# Classes
class BlockPair():
"""Object for tracking source to dest recovery data."""
def __init__(self, source, destination, model, working_dir):
"""Initialize BlockPair()
NOTE: source should be a wk.hw.obj.Disk() object
and destination should be a pathlib.Path() object.
"""
self.sector_size = source.phy_sec
self.source = source.path
self.destination = destination
self.map_data = {}
self.map_path = None
self.size = source.size
self.status = {
def __init__(
self,
source_dev: hw_disk.Disk,
destination: pathlib.Path,
working_dir: pathlib.Path,
):
self.sector_size: int = source_dev.phy_sec
self.source: pathlib.Path = pathlib.Path(source_dev.path)
self.destination: pathlib.Path = destination
self.map_data: dict[str, bool | int] = {}
self.map_path: pathlib.Path = pathlib.Path()
self.size: int = source_dev.size
self.status: dict[str, float | int | str] = {
'read-skip': 'Pending',
'read-full': 'Pending',
'trim': 'Pending',
'scrape': 'Pending',
}
self.view_map = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ
self.view_proc = None
self.view_map: bool = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ
self.view_proc: subprocess.Popen | None = None
# Set map path
# e.g. '(Clone|Image)_Model[_p#]_Size[_Label].map'
map_name = model if model else 'None'
if source.bus == 'Image':
map_name = source_dev.model
if source_dev.bus == 'Image':
map_name = 'Image'
if source.parent:
part_num = re.sub(r"^.*?(\d+)$", r"\1", source.path.name)
if source_dev.parent:
part_num = re.sub(r"^.*?(\d+)$", r"\1", self.source.name)
map_name += f'_p{part_num}'
size_str = std.bytes_to_string(
size=self.size,
use_binary=False,
)
map_name += f'_{size_str.replace(" ", "")}'
if source.raw_details.get('label', ''):
map_name += f'_{source.raw_details["label"]}'
if source_dev.raw_details.get('label', ''):
map_name += f'_{source_dev.raw_details["label"]}'
map_name = map_name.replace(' ', '_')
map_name = map_name.replace('/', '_')
map_name = map_name.replace('\\', '_')
@ -197,7 +197,7 @@ class BlockPair():
"""Get percent rescued from map_data, returns float."""
return 100 * self.map_data.get('rescued', 0) / self.size
def get_rescued_size(self) -> float | int:
def get_rescued_size(self) -> int:
"""Get rescued size using map data.
NOTE: Returns 0 if no map data is available.
@ -211,7 +211,7 @@ class BlockPair():
NOTE: If the file is missing it is assumed that recovery hasn't
started yet so default values will be returned instead.
"""
data = {'full recovery': False, 'pass completed': False}
data: dict[str, bool | int] = {'full recovery': False, 'pass completed': False}
# Get output from ddrescuelog
cmd = [
@ -353,9 +353,8 @@ class State():
"""Add BlockPair object and run safety checks."""
self.block_pairs.append(
BlockPair(
source=source,
source_dev=source,
destination=destination,
model=self.source.model,
working_dir=self.working_dir,
))