chgrp: Move a bunch of stuff to the core modules

Expertcoderz · Expertcoderz · commit 6383f56375ac · 2025-04-12T11:52:54.000Z
diff --git a/userland/core/__init__.py b/userland/core/__init__.py
@@ -1,3 +1,4 @@
 from .command import *
 from .io import *
+from .paths import *
 from .users import *
diff --git a/userland/core/paths.py b/userland/core/paths.py
@@ -0,0 +1,34 @@
+import sys
+from pathlib import Path
+from typing import Generator, Literal
+
+
+def traverse_files(
+    filenames: list[str],
+    recurse_mode: Literal["L", "H", "P"] | None = None,
+    preserve_root: bool = False,
+) -> Generator[Path | None]:
+    if not recurse_mode:
+        yield from map(Path, filenames)
+        return
+
+    def traverse(file: Path) -> Generator[Path]:
+        for child in file.iterdir():
+            if child.is_dir(follow_symlinks=recurse_mode == "L"):
+                yield from traverse(child)
+            yield child
+
+    for file in map(Path, filenames):
+        if preserve_root and file.root == str(file):
+            print(
+                f"recursive operation on '{file}' prevented;"
+                " use --no-preserve-root to override",
+                file=sys.stderr,
+            )
+            yield None
+            continue
+
+        if file.is_dir(follow_symlinks=recurse_mode in set("HL")):
+            yield from traverse(file)
+        else:
+            yield file
diff --git a/userland/core/users.py b/userland/core/users.py
@@ -1,3 +1,4 @@
+import functools
 import grp
 import pwd
 
@@ -8,7 +9,7 @@ def parse_onwer_spec(
     parser: OptionParser, owner_spec: str
 ) -> tuple[int | None, int | None]:
     """
-    Process a string in the form ``[USER][:GROUP]`` and return the UID and GID.
+    Accept a string in the form ``[USER][:[GROUP]]`` and return the UID and GID.
     Either or both may be None if omitted from the input string.
     An appropriate parser error is thrown if obtaining the UID or GID fails.
     """
@@ -18,21 +19,55 @@ def parse_onwer_spec(
     gid: int | None = None
 
     if tokens[0]:
-        if tokens[0].isdecimal():
-            uid = int(tokens[0])
-        else:
-            try:
-                uid = pwd.getpwnam(tokens[0])
-            except KeyError:
-                parser.error(f"invalid user: '{tokens}'")
+        uid = parse_user(parser, tokens[0])
 
     if len(tokens) > 1 and tokens[1]:
-        if tokens[1].isdecimal():
-            gid = int(tokens[1])
-        else:
-            try:
-                gid = grp.getgrnam(tokens[1])
-            except KeyError:
-                parser.error(f"invalid group: '{tokens}'")
+        gid = parse_group(parser, tokens[1])
 
     return uid, gid
+
+
+@functools.lru_cache(1000)
+def parse_user(parser: OptionParser, user: str) -> int:
+    """
+    Accept a string representing a username or UID and return the UID.
+    An appropriate parser error is thrown if obtaining the UID fails.
+    """
+    if user.isdecimal():
+        return int(user)
+
+    try:
+        return pwd.getpwnam(user).pw_uid
+    except KeyError:
+        parser.error(f"invalid user: {user}")
+
+
+@functools.lru_cache(1000)
+def parse_group(parser: OptionParser, group: str) -> int:
+    """
+    Accept a string representing a group name or GID and return the GID.
+    An appropriate parser error is thrown if obtaining the GID fails.
+    """
+    if group.isdecimal():
+        return int(group)
+
+    try:
+        return grp.getgrnam(group).gr_gid
+    except KeyError:
+        parser.error(f"invalid group: {group}")
+
+
+@functools.lru_cache(1000)
+def user_display_name_from_id(uid: int) -> str:
+    try:
+        return pwd.getpwuid(uid).pw_name
+    except KeyError:
+        return str(uid)
+
+
+@functools.lru_cache(1000)
+def group_display_name_from_id(gid: int) -> str:
+    try:
+        return grp.getgrgid(gid).gr_name
+    except KeyError:
+        return str(gid)
diff --git a/userland/utilities/chgrp.py b/userland/utilities/chgrp.py
@@ -1,5 +1,3 @@
-import grp
-import pwd
 import shutil
 import sys
 from pathlib import Path
@@ -152,18 +150,22 @@ def python_userland_chgrp(opts, args):
         if not args:
             parser.error(f"missing operand after '{gname}'")
 
-        if gname.isdecimal():
-            gid = int(gname)
-        else:
-            try:
-                gid = grp.getgrnam(gname).gr_gid
-            except KeyError:
-                parser.error(f"invalid group: '{gname}'")
+        gid = core.parse_group(parser, gname)
 
     failed = False
 
-    def chown(file: Path) -> None:
-        nonlocal failed
+    for file in core.traverse_files(
+        (
+            tqdm(args, ascii=True, desc="Changing group ownership")
+            if opts.progress
+            else args
+        ),
+        recurse_mode=opts.recurse_mode if opts.recursive else None,
+        preserve_root=opts.preserve_root,
+    ):
+        if not file:
+            failed = True
+            continue
 
         try:
             stat = file.stat(follow_symlinks=opts.dereference)
@@ -177,12 +179,9 @@ def chown(file: Path) -> None:
                     f"failed to change group of '{file}' to {gname or gid}",
                     file=sys.stderr,
                 )
-            return
+            continue
 
-        try:
-            prev_gname = grp.getgrgid(prev_gid).gr_name
-        except KeyError:
-            prev_gname = str(prev_gid)
+        prev_gname = core.group_display_name_from_id(prev_gid)
 
         # Note: while it's possible, we do not check if prev_gid == gid at
         # this point because even if they are the same, an error should be
@@ -193,7 +192,7 @@ def chown(file: Path) -> None:
         ):
             if opts.verbosity > 2:
                 print(f"group of '{file}' retained as {prev_gname}")
-            return
+            continue
 
         try:
             shutil.chown(file, group=gid, follow_symlinks=opts.dereference)
@@ -206,49 +205,12 @@ def chown(file: Path) -> None:
                         f"failed to change group of '{file}' to {gname or gid}",
                         file=sys.stderr,
                     )
-            return
+            continue
 
         if prev_gid == gid:
             if opts.verbosity > 2:
                 print(f"group of '{file}' retained as {prev_gname}")
         elif opts.verbosity > 1:
             print(f"changed group of '{file}' from {prev_gname} to {gname or gid}")
 
-    files = map(
-        Path,
-        (
-            tqdm(args, ascii=True, desc="Changing group ownership")
-            if opts.progress
-            else args
-        ),
-    )
-
-    if opts.recursive:
-
-        def traverse(file: Path) -> None:
-            for child in file.iterdir():
-                if child.is_dir(follow_symlinks=opts.recurse_mode == "L"):
-                    traverse(child)
-                chown(file)
-
-        for file in files:
-            if file.is_dir(
-                follow_symlinks=opts.recurse_mode == "H" or opts.recurse_mode == "L"
-            ):
-                if opts.preserve_root and file.root == str(file):
-                    failed = True
-                    print(
-                        f"recursive operation on '{file}' prevented;"
-                        " use --no-preserve-root to override",
-                        file=sys.stderr,
-                    )
-                    continue
-
-                traverse(file)
-            else:
-                chown(file)
-    else:
-        for file in files:
-            chown(file)
-
     return int(failed)
