diff --git a/src/install/loadingbar.py b/src/install/loadingbar.py index ed7419c..34d1e84 100644 --- a/src/install/loadingbar.py +++ b/src/install/loadingbar.py @@ -15,8 +15,13 @@ # along with this program. If not, see . from os import get_terminal_size -from typing import Collection, Generic, Literal, TypeVar, Optional, overload, Type -from types import TracebackType +from typing import ( + Collection, Generic, Literal, TypeVar, + Optional, overload, TYPE_CHECKING +) + +if TYPE_CHECKING: + from types import TracebackType from .filesize import size, traditional @@ -95,9 +100,6 @@ def __init__( :param disappear: If you want the bar to disappear """ - # Define class variables - self._iter_type: Literal['iter', 'total', 'both'] - if iterator is not None and total == 0: # iter self.iterator = iterator @@ -165,10 +167,14 @@ def __next__(self) -> _T: if self.iterable is not None: self.item = next(self.iterable) self.idx += 1 - elif type(self.total) == int and self.idx < self.total: # total - self.idx += 1 else: - raise StopIteration + if TYPE_CHECKING and type(self.total) != int: + raise TypeError + + if self.idx < self.total: # total + self.idx += 1 + else: + raise StopIteration # Refresh the loading bar self.refresh() @@ -184,9 +190,9 @@ def __enter__(self) -> "loadingbar[_T]": return self - def __exit__(self, exc_type: Optional[Type[BaseException]], + def __exit__(self, exc_type: Optional[type[BaseException]], exc_value: Optional[BaseException], - traceback: Optional[TracebackType]) -> bool: + traceback: Optional['TracebackType']) -> bool: "Allow a with-as statement" return False # No errors @@ -195,7 +201,10 @@ def refresh(self) -> None: "Refresh the loading bar, called automatically" # Calculate progress - if type(self.total) == int and self.total != 0: # both, total + if hasattr(self, 'total'): # both, total + if TYPE_CHECKING and type(self.total) != int: + raise TypeError + percent = round(self.idx / self.total * 100, 0) else: percent = round(self.idx / self.iterator_len * 100, 0) @@ -206,12 +215,14 @@ def refresh(self) -> None: # Define the current and total if self.unit == 'it': current = str(self.idx) - total = str(self.total if self.total != 0 else self.iterator_len) + total = str(self.total if hasattr( + self, 'total') else self.iterator_len) else: current = size(self.idx, traditional) - if type(arg := self.total if self.total != 0 - else self.iterator_len) != int: - raise TypeError # Shouldn't be able to happen + arg = self.total if hasattr(self, 'total') else self.iterator_len + + if TYPE_CHECKING and type(arg) != int: + raise TypeError total = size(arg, traditional) @@ -240,8 +251,7 @@ def refresh(self) -> None: ), end=end) # Clear the loading bar at the end - if self.idx == (self.total if self._iter_type in ( - 'both', 'total') else self.iterator_len): + if self.idx == (self.total if hasattr(self, 'total') else self.iterator_len): if self.disappear and self.show_desc: print('\r\033[K\033[F\r\033[K', end='') elif self.disappear: @@ -251,7 +261,7 @@ def refresh(self) -> None: def update(self, amount: int) -> None: "Add 'n' amount of iterations to the loading bar" - if self._iter_type == 'iter': + if not hasattr(self, 'total'): i = 0 # Call next(self) while less than amount