Select Counts With Command-Line Options (Solution: Part 2)

In this lesson, you’ll finish off your implementation of the task from part one by incorporating the logic to selectively display the requested counts.

Know Your Starting Point

If you got lost while working on this task, then expand the section below to view the full source code of the wordcount command from part one:

Pythonsrc/wordcount.py
      
    
import sys
from argparse import ArgumentParser, Namespace
from dataclasses import dataclass
from enum import IntFlag, auto
from functools import cached_property
from pathlib import Path
from typing import NamedTuple

class SelectedCounts(IntFlag):
    NONE = 0
    LINES = auto()
    WORDS = auto()
    BYTES = auto()
    DEFAULT = LINES | WORDS | BYTES

class Arguments(Namespace):
    @cached_property
    def selected_counts(self):
        selected = self.lines | self.words | self.bytes
        return selected or SelectedCounts.DEFAULT

class Counts(NamedTuple):
    lines: int = 0
    words: int = 0
    bytes: int = 0

    @property
    def max_digits(self):
        return len(str(max(self)))

    def as_string(self, max_digits):
        return (
            f"{self.lines:>{max_digits}} "
            f"{self.words:>{max_digits}} "
            f"{self.bytes:>{max_digits}}"
        )

    def __add__(self, other):
        return Counts(
            lines=self.lines + other.lines,
            words=self.words + other.words,
            bytes=self.bytes + other.bytes,
        )

@dataclass(frozen=True)
class FileInfo:
    path: Path
    counts: Counts

    @classmethod
    def from_path(cls, path):
        if path.name == "-":
            raw_text = sys.stdin.buffer.read()
        elif path.is_file():
            raw_text = path.read_bytes()
        else:
            return cls(path, Counts())
        text = raw_text.decode("utf-8")
        return cls(
            path,
            Counts(
                lines=text.count("\n"),
                words=len(text.split()),
                bytes=len(raw_text),
            ),
        )

def main():
    args = parse_args()
    if len(args.paths) > 0:
        file_infos = [FileInfo.from_path(path) for path in args.paths]
    else:
        file_infos = [FileInfo.from_path(Path("-"))]
    total_counts = sum((info.counts for info in file_infos), Counts())
    max_digits = total_counts.max_digits
    for info in file_infos:
        line = info.counts.as_string(max_digits)
        if info.path == Path("-"):
            print(line)
        elif not info.path.exists():
            print(line, info.path, "(no such file or directory)")
        elif info.path.is_dir():
            print(line, f"{info.path}/ (is a directory)")
        else:
            print(line, info.path)
    if len(file_infos) > 1:
        print(total_counts.as_string(max_digits), "total")

def parse_args():
    parser = ArgumentParser()
    parser.add_argument("paths", nargs="*", type=Path)
    for flag in SelectedCounts:
        parser.add_argument(
            f"--{flag.name.lower()}",
            action="store_const",
            const=flag,
            default=SelectedCounts.NONE,
        )
    return parser.parse_args(namespace=Arguments())

While this code makes a few more acceptance criteria pass, most still keep failing. Don’t worry, though, as you’re on the right track. You’ve already done the hard work in part one!

Locked learning resources

Join us and get access to thousands of tutorials and a community of expert Pythonistas.

Unlock This Lesson

Already a member? Sign-In

Locked learning resources

The full lesson is for members only. Join us and get access to thousands of tutorials and a community of expert Pythonistas.

Unlock This Lesson

Already a member? Sign-In

davidbonn on June 20, 2025

Maybe it is just me but it seemed to be simpler to have the command line options represented as a list of strs, the strs being the attribute names I wanted to print, e.g. you could do:

    ap.add_argument("--lines", action="append_const", const="n_lines", dest="options")
    ap.add_argument("--words", action="append_const", const="n_words", dest="options")
    ap.add_argument("--bytes", action="append_const", const="n_bytes", dest="options")
    ap.add_argument("--chars", action="append_const", const="n_chars", dest="options")

And the code to generate the output would be approximately:

    def output(self, options, width):
        values = []
        for thing in ["n_lines", "n_words", "n_chars", "n_bytes"]:
            if thing in options:
                values.append(getattr(self, thing))

        rc = " ".join([f"{v:>{width}d}" for v in values])

        if len(self.filename):
            rc += f" {self.filename}{self.message}"

        return rc

Become a Member to join the conversation.