Skip to content

Commit

Permalink
Fixed several typing issues and hopefully some more windows bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
wolph committed Feb 23, 2024
1 parent afc42ff commit 0046fdf
Show file tree
Hide file tree
Showing 6 changed files with 90 additions and 57 deletions.
2 changes: 1 addition & 1 deletion examples.py
Expand Up @@ -61,7 +61,7 @@ def do_something(bar):
# Sleep up to 0.1 seconds
time.sleep(random.random() * 0.1)

with (progressbar.MultiBar() as multibar):
with progressbar.MultiBar() as multibar:
bar_labels = []
for i in range(BARS):
# Get a progressbar
Expand Down
57 changes: 36 additions & 21 deletions progressbar/__main__.py
Expand Up @@ -4,7 +4,9 @@
import contextlib
import pathlib
import sys
from typing import BinaryIO
import typing
from pathlib import Path
from typing import BinaryIO, TextIO

import progressbar

Expand Down Expand Up @@ -52,7 +54,7 @@ def size_to_bytes(size_str: str) -> int:
size_str = size_str[:-1]

# Convert the size_str to an integer and apply the exponent
return int(size_str) * (1024 ** exponent)
return int(size_str) * (1024**exponent)


def create_argument_parser() -> argparse.ArgumentParser:
Expand All @@ -63,7 +65,7 @@ def create_argument_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description='''
Monitor the progress of data through a pipe.
Note that this is a Python implementation of the original `pv` command
that is functional but not yet feature complete.
'''
Expand Down Expand Up @@ -270,28 +272,26 @@ def create_argument_parser() -> argparse.ArgumentParser:
def main(argv: list[str] | None = None): # noqa: C901
'''
Main function for the `progressbar` command.
'''
parser = create_argument_parser()
args = parser.parse_args(argv)
binary_mode = '' if args.line_mode else 'b'
Args:
argv (list[str] | None): Command-line arguments passed to the script.
Returns:
None
'''
parser: argparse.ArgumentParser = create_argument_parser()
args: argparse.Namespace = parser.parse_args(argv)

with contextlib.ExitStack() as stack:
if args.output and args.output != '-':
output_stream = stack.enter_context(
open(args.output, 'w' + binary_mode)
)
else:
if args.line_mode:
output_stream = sys.stdout
else:
output_stream = sys.stdout.buffer
output_stream: typing.IO[typing.Any] = _get_output_stream(
args.output, args.line_mode, stack
)

input_paths = []
total_size = 0
filesize_available = True
input_paths: list[BinaryIO | TextIO | Path] = []
total_size: int = 0
filesize_available: bool = True
for filename in args.input:
input_path: BinaryIO | pathlib.Path
input_path: typing.IO[typing.Any] | pathlib.Path
if filename == '-':
if args.line_mode:
input_path = sys.stdin
Expand Down Expand Up @@ -356,12 +356,13 @@ def main(argv: list[str] | None = None): # noqa: C901
for input_path in input_paths:
if isinstance(input_path, pathlib.Path):
input_stream = stack.enter_context(
input_path.open('r' + binary_mode)
input_path.open('r' if args.line_mode else 'rb')
)
else:
input_stream = input_path

while True:
data: str | bytes
if args.line_mode:
data = input_stream.readline(buffer_size)
else:
Expand All @@ -377,5 +378,19 @@ def main(argv: list[str] | None = None): # noqa: C901
bar.finish(dirty=True)


def _get_output_stream(
output: str | None,
line_mode: bool,
stack: contextlib.ExitStack,
) -> typing.IO[typing.Any]:
if output and output != '-':
mode = 'w' if line_mode else 'wb'
return stack.enter_context(open(output, mode)) # noqa: SIM115
elif line_mode:
return sys.stdout
else:
return sys.stdout.buffer


if __name__ == '__main__':
main()
70 changes: 44 additions & 26 deletions progressbar/bar.py
Expand Up @@ -162,13 +162,13 @@ class DefaultFdMixin(ProgressBarMixinBase):
#: Set the terminal to be ANSI compatible. If a terminal is ANSI
#: compatible we will automatically enable `colors` and disable
#: `line_breaks`.
is_ansi_terminal: bool = False
is_ansi_terminal: bool | None = False
#: Whether the file descriptor is a terminal or not. This is used to
#: determine whether to use ANSI escape codes or not.
is_terminal: bool
is_terminal: bool | None
#: Whether to print line breaks. This is useful for logging the
#: progressbar. When disabled the current line is overwritten.
line_breaks: bool = True
line_breaks: bool | None = True
#: Specify the type and number of colors to support. Defaults to auto
#: detection based on the file descriptor type (i.e. interactive terminal)
#: environment variables such as `COLORTERM` and `TERM`. Color output can
Expand All @@ -179,9 +179,7 @@ class DefaultFdMixin(ProgressBarMixinBase):
#: For true (24 bit/16M) color support you can use `COLORTERM=truecolor`.
#: For 256 color support you can use `TERM=xterm-256color`.
#: For 16 colorsupport you can use `TERM=xterm`.
enable_colors: progressbar.env.ColorSupport | bool | None = (
progressbar.env.COLOR_SUPPORT
)
enable_colors: progressbar.env.ColorSupport = progressbar.env.COLOR_SUPPORT

def __init__(
self,
Expand All @@ -200,7 +198,7 @@ def __init__(
fd = self._apply_line_offset(fd, line_offset)
self.fd = fd
self.is_ansi_terminal = progressbar.env.is_ansi_terminal(fd)
self.is_terminal = self._determine_is_terminal(fd, is_terminal)
self.is_terminal = progressbar.env.is_terminal(fd, is_terminal)
self.line_breaks = self._determine_line_breaks(line_breaks)
self.enable_colors = self._determine_enable_colors(enable_colors)

Expand All @@ -219,29 +217,47 @@ def _apply_line_offset(
else:
return fd

def _determine_is_terminal(
self,
fd: base.TextIO,
is_terminal: bool | None,
) -> bool:
if is_terminal is not None:
return progressbar.env.is_terminal(fd, is_terminal)
else:
return progressbar.env.is_ansi_terminal(fd)

def _determine_line_breaks(self, line_breaks: bool | None) -> bool:
def _determine_line_breaks(self, line_breaks: bool | None) -> bool | None:
if line_breaks is None:
return progressbar.env.env_flag(
'PROGRESSBAR_LINE_BREAKS',
not self.is_terminal,
)
else:
return bool(line_breaks)
return line_breaks

def _determine_enable_colors(
self,
enable_colors: progressbar.env.ColorSupport | None,
) -> progressbar.env.ColorSupport:
'''
Determines the color support for the progress bar.
This method checks the `enable_colors` parameter and the environment
variables `PROGRESSBAR_ENABLE_COLORS` and `FORCE_COLOR` to determine
the color support.
If `enable_colors` is:
- `None`, it checks the environment variables and the terminal
compatibility to ANSI.
- `True`, it sets the color support to XTERM_256.
- `False`, it sets the color support to NONE.
- For different values that are not instances of
`progressbar.env.ColorSupport`, it raises a ValueError.
Args:
enable_colors (progressbar.env.ColorSupport | None): The color
support setting from the user. It can be None, True, False,
or an instance of `progressbar.env.ColorSupport`.
Returns:
progressbar.env.ColorSupport: The determined color support.
Raises:
ValueError: If `enable_colors` is not None, True, False, or an
instance of `progressbar.env.ColorSupport`.
'''
color_support = progressbar.env.ColorSupport.NONE
if enable_colors is None:
colors = (
progressbar.env.env_flag('PROGRESSBAR_ENABLE_COLORS'),
Expand All @@ -252,21 +268,23 @@ def _determine_enable_colors(
for color_enabled in colors:
if color_enabled is not None:
if color_enabled:
enable = progressbar.env.COLOR_SUPPORT
color_support = progressbar.env.COLOR_SUPPORT
else:
enable = progressbar.env.ColorSupport.NONE
color_support = progressbar.env.ColorSupport.NONE
break
else:
enable = False
color_support = progressbar.env.ColorSupport.NONE

elif enable_colors is True:
enable = progressbar.env.ColorSupport.XTERM_256
color_support = progressbar.env.ColorSupport.XTERM_256
elif enable_colors is False:
enable = progressbar.env.ColorSupport.NONE
elif not isinstance(enable_colors, progressbar.env.ColorSupport):
color_support = progressbar.env.ColorSupport.NONE
elif isinstance(enable_colors, progressbar.env.ColorSupport):
color_support = enable_colors
else:
raise ValueError(f'Invalid color support value: {enable_colors}')

return enable
return color_support

def print(self, *args: types.Any, **kwargs: types.Any) -> None:
print(*args, file=self.fd, **kwargs)
Expand Down
6 changes: 3 additions & 3 deletions progressbar/env.py
Expand Up @@ -81,7 +81,7 @@ def from_env(cls):
):
return cls.XTERM_TRUECOLOR
else:
return cls.WINDOWS
return cls.WINDOWS # pragma: no cover

support = cls.NONE
for variable in variables:
Expand Down Expand Up @@ -142,7 +142,7 @@ def is_ansi_terminal(
return is_terminal


def is_terminal(fd: base.IO, is_terminal: bool | None = None) -> bool:
def is_terminal(fd: base.IO, is_terminal: bool | None = None) -> bool | None:
if is_terminal is None:
# Full ansi support encompasses what we expect from a terminal
is_terminal = is_ansi_terminal(fd) or None
Expand All @@ -159,7 +159,7 @@ def is_terminal(fd: base.IO, is_terminal: bool | None = None) -> bool:
except Exception:
is_terminal = False

return bool(is_terminal)
return is_terminal


# Enable Windows full color mode if possible
Expand Down
8 changes: 4 additions & 4 deletions progressbar/terminal/os_specific/__init__.py
Expand Up @@ -11,13 +11,13 @@
else:
from .posix import getch as _getch

def _reset_console_mode():
def _reset_console_mode() -> None:
pass

def _set_console_mode():
pass
def _set_console_mode() -> bool:
return False

def _get_console_mode():
def _get_console_mode() -> int:
return 0


Expand Down
4 changes: 2 additions & 2 deletions progressbar/terminal/os_specific/windows.py
Expand Up @@ -120,7 +120,7 @@ class _Event(ctypes.Union):
_fields_ = (('EventType', _WORD), ('Event', _Event))


def reset_console_mode():
def reset_console_mode() -> None:
_SetConsoleMode(_HANDLE(_h_console_input), _DWORD(_input_mode.value))
_SetConsoleMode(_HANDLE(_h_console_output), _DWORD(_output_mode.value))

Expand All @@ -144,7 +144,7 @@ def get_console_mode() -> int:
return _input_mode.value


def set_text_color(color):
def set_text_color(color) -> None:
_kernel32.SetConsoleTextAttribute(_h_console_output, color)


Expand Down

0 comments on commit 0046fdf

Please sign in to comment.