Add wc

For counting

Author: Expertcoderz

userland/utilities/wc.py

@@ -0,0 +1,242 @@
+import itertools
+import sys
+from collections.abc import Callable, Iterable, Sequence
+from enum import Enum, auto
+from typing import Self, TextIO
+
+from .. import core
+
+
+def wc_count_bytes(stream: Iterable[str]) -> int:
+    return sum(len(line.encode()) for line in stream)
+
+
+def wc_count_bytes_fast(stream: Iterable[bytes]) -> int:
+    return sum(len(line) for line in stream)
+
+
+def wc_count_chars(stream: Iterable[str]) -> int:
+    return sum(len(line) for line in stream)
+
+
+def wc_count_words(stream: Iterable[str]) -> int:
+    return sum(len(line.split()) for line in stream)
+
+
+def wc_count_lines(stream: Iterable[bytes]) -> int:
+    return sum(1 for line in stream if line.endswith("\n"))
+
+
+def wc_get_max_line_len(stream: Iterable[bytes]) -> int:
+    return max(map(len, stream))
+
+
+wc_enabled_funcs: dict[Callable[[Iterable[bytes]], int], bool] = dict(
+    (func, False)
+    for func in (
+        wc_count_lines,
+        wc_count_words,
+        wc_count_chars,
+        wc_count_bytes,
+        wc_get_max_line_len,
+    )
+)
+
+
+class TotalMode(Enum):
+    AUTO = auto()
+    ALWAYS = auto()
+    ONLY = auto()
+    NEVER = auto()
+
+    @classmethod
+    def from_string(cls, string: str | None) -> Self:
+        match string:
+            case None | "auto":
+                return cls.AUTO
+            case "always":
+                return cls.ALWAYS
+            case "only":
+                return cls.ONLY
+            case "never":
+                return cls.NEVER
+
+        raise ValueError(f"invalid value '{string}' for --total")
+
+
+def format_counts(counts: Sequence[int]) -> str:
+    return " ".join(map(str, counts))
+
+
+def display_lines(
+    lines: Iterable[tuple[str, Sequence[int]] | tuple[None, None]],
+    total_mode: TotalMode,
+    n_counts: int,
+) -> bool:
+    n_lines = 0
+    total = [0] * n_counts
+
+    failed = False
+
+    for filename, counts in lines:
+        n_lines += 1
+
+        if filename is None:
+            failed = True
+            continue
+
+        if total_mode != TotalMode.ONLY:
+            print(
+                f"{format_counts(counts)}"
+                if filename is None
+                else f"{format_counts(counts)} {filename}"
+            )
+
+        for i, count in enumerate(counts):
+            total[i] += count
+
+    if total_mode == TotalMode.ONLY:
+        print(format_counts(total))
+    elif total_mode == TotalMode.ALWAYS or (
+        total_mode == TotalMode.AUTO and n_lines > 1
+    ):
+        print(f"{format_counts(total)} total")
+
+    return failed
+
+
+def count_file(filename: str | None) -> tuple[str, Sequence[int]] | tuple[None, None]:
+    if filename == "":
+        core.perror("empty filename")
+        return None, None
+
+    bytes_only = (
+        wc_count_bytes in wc_enabled_funcs
+        and sum(1 for enabled in wc_enabled_funcs.values() if enabled) < 2
+    )
+
+    if filename is None or filename == "-":
+        return filename, (
+            (wc_count_bytes_fast(sys.stdin.buffer),)
+            if bytes_only
+            else wc_apply_funcs(sys.stdin)
+        )
+
+    try:
+        if bytes_only:
+            with open(filename, "rb") as f:
+                return filename, (wc_count_bytes_fast(f),)
+        else:
+            with open(filename, "r", encoding="utf-8") as f:
+                return filename, wc_apply_funcs(f)
+    except OSError as e:
+        core.perror(e)
+        return None, None
+
+
+def wc_apply_funcs(stream: Iterable[str]) -> Sequence[int]:
+    streams = itertools.tee(stream, len(wc_enabled_funcs))
+    counts = []
+
+    for i, (wc_func, enabled) in enumerate(wc_enabled_funcs.items()):
+        if enabled:
+            counts.append(wc_func(streams[i]))
+
+    return counts
+
+
+def wc_get_enabler(func) -> Callable:
+    def callback(*_):
+        wc_enabled_funcs[func] = True
+
+    return callback
+
+
+def get_filenames(source_file: TextIO) -> Iterable[str]:
+    # Note that coreutils wc reads all the names from
+    # the source file before starting to count anything.
+    filenames = source_file.read().split("\0")
+    if not filenames[-1]:
+        filenames = filenames[:-1]
+
+    return filenames
+
+
+parser = core.ExtendedOptionParser(
+    usage="%prog [FILE]...", description="Print newline, word and byte counts."
+)
+
+parser.add_option(
+    "-l",
+    "--lines",
+    action="callback",
+    callback=wc_get_enabler(wc_count_lines),
+    help="print line counts",
+)
+parser.add_option(
+    "-w",
+    "--words",
+    action="callback",
+    callback=wc_get_enabler(wc_count_words),
+    help="print word counts",
+)
+parser.add_option(
+    "-m",
+    "--chars",
+    action="callback",
+    callback=wc_get_enabler(wc_count_chars),
+    help="print character counts",
+)
+parser.add_option(
+    "-c",
+    "--bytes",
+    action="callback",
+    callback=wc_get_enabler(wc_count_bytes),
+    help="print byte counts",
+)
+parser.add_option(
+    "-L",
+    "--max-line-length",
+    action="callback",
+    callback=wc_get_enabler(wc_get_max_line_len),
+    help="print length of longest line",
+)
+parser.add_option(
+    "--total",
+    action="store",
+    dest="total",
+    metavar="WHEN",
+    help="set whether to print total counts",
+)
+parser.add_option(
+    "--files0-from", action="store", dest="filename_source", metavar="FILE", help="read"
+)
+
+
+@core.command(parser)
+def python_userland_wc(opts, args: list[str]) -> int:
+    if not any(wc_enabled_funcs.values()):
+        wc_enabled_funcs[wc_count_bytes] = True
+        wc_enabled_funcs[wc_count_words] = True
+        wc_enabled_funcs[wc_count_lines] = True
+
+    n_cols = sum(1 for enabled in wc_enabled_funcs.values() if enabled)
+
+    try:
+        total_mode = TotalMode.from_string(opts.total)
+    except ValueError as e:
+        core.perror(e)
+        return 1
+
+    if opts.filename_source is not None:
+        parser.expect_nargs(args, 0)
+
+        if opts.filename_source == "-":
+            args = get_filenames(sys.stdin)
+        else:
+            with open(opts.filename_source, "r", encoding="utf-8") as f:
+                args = get_filenames(f)
+    elif not args:
+        args = [None]
+
+    return int(display_lines(map(count_file, args), total_mode, n_cols))