Skip to content

Commit

Permalink
Added progressbar command for commandline progressbars
Browse files Browse the repository at this point in the history
  • Loading branch information
wolph committed Feb 18, 2024
1 parent 671b723 commit 5707548
Show file tree
Hide file tree
Showing 4 changed files with 396 additions and 2 deletions.
7 changes: 7 additions & 0 deletions docs/progressbar.algorithms.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
progressbar.algorithms module
=============================

.. automodule:: progressbar.algorithms
:members:
:undoc-members:
:show-inheritance:
279 changes: 279 additions & 0 deletions progressbar/__main__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,279 @@
import argparse
import contextlib
import pathlib
import sys
import time
from typing import BinaryIO

import progressbar


def size_to_bytes(size_str: str) -> int:
'''
Convert a size string with suffixes 'k', 'm', etc., to bytes.
Note: This function also supports '@' as a prefix to a file path to get the
file size.
>>> size_to_bytes('1024k')
1048576
>>> size_to_bytes('1024m')
1073741824
>>> size_to_bytes('1024g')
1099511627776
>>> size_to_bytes('1024')
1024
>>> size_to_bytes('1024p')
1125899906842624
'''

# Define conversion rates
suffix_exponent = {
'k': 1,
'm': 2,
'g': 3,
't': 4,
'p': 5,
}

# Initialize the default exponent to 0 (for bytes)
exponent = 0

# Check if the size starts with '@' (for file sizes, not handled here)
if size_str.startswith('@'):
return pathlib.Path(size_str[1:]).stat().st_size

# Check if the last character is a known suffix and adjust the multiplier
if size_str[-1].lower() in suffix_exponent:
# Update exponent based on the suffix
exponent = suffix_exponent[size_str[-1].lower()]
# Remove the suffix from the size_str
size_str = size_str[:-1]

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


def create_argument_parser() -> argparse.ArgumentParser:
'''
Create the argument parser for the `progressbar` command.
>>> parser = create_argument_parser()
>>> parser.parse_args(['-p', '-t', '-e', '-r', '-a', '-b', '-8', '-T', '-A', '-F', '-n', '-q', 'input', '-o', 'output'])
Namespace(average_rate=True, bytes=True, eta=True, fineta=False, format=None, height=None, input=['input'], interval=None, last_written=None, line_mode=False, name=None, numeric=True, output='output', progress=True, quiet=True, rate=True, rate_limit=None, remote=None, size=None, stop_at_size=False, sync=False, timer=True, wait=False, watchfd=None, width=None)
Returns:
argparse.ArgumentParser: The argument parser for the `progressbar` command.
'''

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.
''')

# Display switches
parser.add_argument('-p', '--progress', action='store_true',
help='Turn the progress bar on.')
parser.add_argument('-t', '--timer', action='store_true',
help='Turn the timer on.')
parser.add_argument('-e', '--eta', action='store_true',
help='Turn the ETA timer on.')
parser.add_argument('-I', '--fineta', action='store_true',
help='Display the ETA as local time of arrival.')
parser.add_argument('-r', '--rate', action='store_true',
help='Turn the rate counter on.')
parser.add_argument('-a', '--average-rate', action='store_true',
help='Turn the average rate counter on.')
parser.add_argument('-b', '--bytes', action='store_true',
help='Turn the total byte counter on.')
parser.add_argument('-8', '--bits', action='store_true',
help='Display total bits instead of bytes.')
parser.add_argument('-T', '--buffer-percent', action='store_true',
help='Turn on the transfer buffer percentage display.')
parser.add_argument('-A', '--last-written', type=int,
help='Show the last NUM bytes written.')
parser.add_argument('-F', '--format', type=str,
help='Use the format string FORMAT for output format.')
parser.add_argument('-n', '--numeric', action='store_true',
help='Numeric output.')
parser.add_argument('-q', '--quiet', action='store_true', help='No output.')

# Output modifiers
parser.add_argument('-W', '--wait', action='store_true',
help='Wait until the first byte has been transferred.')
parser.add_argument('-D', '--delay-start', type=float, help='Delay start.')
parser.add_argument('-s', '--size', type=str,
help='Assume total data size is SIZE.')
parser.add_argument('-l', '--line-mode', action='store_true',
help='Count lines instead of bytes.')
parser.add_argument('-0', '--null', action='store_true',
help='Count lines terminated with a zero byte.')
parser.add_argument('-i', '--interval', type=float,
help='Interval between updates.')
parser.add_argument('-m', '--average-rate-window', type=int,
help='Window for average rate calculation.')
parser.add_argument('-w', '--width', type=int,
help='Assume terminal is WIDTH characters wide.')
parser.add_argument('-H', '--height', type=int,
help='Assume terminal is HEIGHT rows high.')
parser.add_argument('-N', '--name', type=str,
help='Prefix output information with NAME.')
parser.add_argument('-f', '--force', action='store_true',
help='Force output.')
parser.add_argument('-c', '--cursor', action='store_true',
help='Use cursor positioning escape sequences.')

# Data transfer modifiers
parser.add_argument('-L', '--rate-limit', type=str,
help='Limit transfer to RATE bytes per second.')
parser.add_argument('-B', '--buffer-size', type=str,
help='Use transfer buffer size of BYTES.')
parser.add_argument('-C', '--no-splice', action='store_true',
help='Never use splice.')
parser.add_argument('-E', '--skip-errors', action='store_true',
help='Ignore read errors.')
parser.add_argument('-Z', '--error-skip-block', type=str,
help='Skip block size when ignoring errors.')
parser.add_argument('-S', '--stop-at-size', action='store_true',
help='Stop transferring after SIZE bytes.')
parser.add_argument('-Y', '--sync', action='store_true',
help='Synchronise buffer caches to disk after writes.')
parser.add_argument('-K', '--direct-io', action='store_true',
help='Set O_DIRECT flag on all inputs/outputs.')
parser.add_argument('-X', '--discard', action='store_true',
help='Discard input data instead of transferring it.')
parser.add_argument('-d', '--watchfd', type=str,
help='Watch file descriptor of process.')
parser.add_argument('-R', '--remote', type=int,
help='Remote control another running instance of pv.')

# General options
parser.add_argument('-P', '--pidfile', type=pathlib.Path,
help='Save process ID in FILE.')
parser.add_argument(
'input',
help='Input file path. Uses stdin if not specified.',
default='-',
nargs='*',
)
parser.add_argument(
'-o',
'--output',
default='-',
help='Output file path. Uses stdout if not specified.')

return parser


def main(argv: list[str] = sys.argv[1:]):
'''
Main function for the `progressbar` command.
'''
parser = create_argument_parser()
args = parser.parse_args(argv)

binary_mode = '' if args.line_mode else 'b'

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

input_paths = []
total_size = 0
filesize_available = True
for filename in args.input:
input_path: BinaryIO | pathlib.Path
if filename == '-':
if args.line_mode:
input_path = sys.stdin
else:
input_path = sys.stdin.buffer

filesize_available = False
else:
input_path = pathlib.Path(filename)
if not input_path.exists():
parser.error(f'File not found: {filename}')

if not args.size:
total_size += input_path.stat().st_size

input_paths.append(input_path)

# Determine the size for the progress bar (if provided)
if args.size:
total_size = size_to_bytes(args.size)
filesize_available = True

if filesize_available:
# Create the progress bar components
widgets = [
progressbar.Percentage(),
' ',
progressbar.Bar(),
' ',
progressbar.Timer(),
' ',
progressbar.FileTransferSpeed(),
]
else:
widgets = [
progressbar.SimpleProgress(),
' ',
progressbar.DataSize(),
' ',
progressbar.Timer(),
]

if args.eta:
widgets.append(' ')
widgets.append(progressbar.AdaptiveETA())

# Initialize the progress bar
bar = progressbar.ProgressBar(
# widgets=widgets,
max_value=total_size or None,
max_error=False,
)

# Data processing and updating the progress bar
buffer_size = size_to_bytes(
args.buffer_size) if args.buffer_size else 1024
total_transferred = 0

bar.start()
with contextlib.suppress(KeyboardInterrupt):
for input_path in input_paths:
if isinstance(input_path, pathlib.Path):
input_stream = stack.enter_context(
input_path.open('r' + binary_mode))
else:
input_stream = input_path

while True:
if args.line_mode:
data = input_stream.readline(buffer_size)
else:
data = input_stream.read(buffer_size)

if not data:
break

output_stream.write(data)
total_transferred += len(data)
bar.update(total_transferred)

bar.finish(dirty=True)


if __name__ == '__main__':
main()
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -108,8 +108,8 @@ exclude = ['docs*', 'tests*']
[tool.setuptools]
include-package-data = true

# [project.scripts]
# progressbar2 = 'progressbar.cli:main'
[project.scripts]
progressbar = 'progressbar.cli:main'

[project.optional-dependencies]
docs = ['sphinx>=1.8.5', 'sphinx-autodoc-typehints>=1.6.0']
Expand Down
Loading

0 comments on commit 5707548

Please sign in to comment.