Move walk function into utils

3 files changed, 17 insertions, 29 deletions

diff --git a/patternutils/commands/pmatch.py b/patternutils/commands/pmatch.py
index 521aa9d..18df269 100644
--- a/patternutils/commands/pmatch.py
+++ b/patternutils/commands/pmatch.py
@@ -1,10 +1,8 @@
 import functools
 import itertools
-import json
-import os
 import os.path
 import re
-from typing import Any, Dict, Generator, Iterator, List, Pattern, Tuple
+from typing import Any, Dict, Iterator, List, Pattern, Tuple
 
 import click
 
@@ -13,19 +11,6 @@ from .. import match as mat
 from .. import utils
 
 
-def walk(
-    dir: str, recursive: bool = False, match_directories: bool = False
-) -> Generator[os.DirEntry, None, None]:
-    for entry in os.scandir(dir):
-        is_dir = entry.is_dir()
-        if recursive and is_dir:
-            yield from walk(entry.path, recursive, match_directories)
-        elif entry.is_file() or (match_directories and is_dir):
-            yield entry
-        else:
-            continue
-
-
 def apply_regex(
     walk_function: Iterator[Any], regex_pattern: Pattern[str], match_full_path: bool,
 ) -> Tuple[List[Dict[str, str]], bool]:
@@ -135,7 +120,7 @@ def pmatch(
 
     # Monkey patch walk function to include commandline options
     walk_function = functools.partial(
-        walk, recursive=recursive, match_directories=match_directories
+        utils.walk, recursive=recursive, match_directories=match_directories
     )
     new_walk = itertools.chain.from_iterable(map(walk_function, directory))
 
diff --git a/patternutils/commands/prss.py b/patternutils/commands/prss.py
index 1aff122..b2a83fc 100644
--- a/patternutils/commands/prss.py
+++ b/patternutils/commands/prss.py
@@ -1,6 +1,4 @@
-import os
 import os.path
-from typing import Generator
 
 import click
 
@@ -12,14 +10,6 @@ except ImportError:
 from .. import utils
 
 
-def walk(dir: str) -> Generator[os.DirEntry, None, None]:
-    for entry in os.scandir(dir):
-        if entry.is_file():
-            yield entry
-        else:
-            continue
-
-
 @click.command(context_settings={"help_option_names": ("-h", "--help", "-?")})
 @click.option("-d", "--directory", default="./")
 @click.option(
@@ -54,7 +44,7 @@ def prss(
     direntries = []
     filenames = []
 
-    for f in walk(directory):
+    for f in utils.walk(directory, recursive=False, match_directories=False):
         direntries.append(f)
         filenames.append(f.name)
 
diff --git a/patternutils/utils.py b/patternutils/utils.py
index b76acee..9160cfe 100644
--- a/patternutils/utils.py
+++ b/patternutils/utils.py
@@ -1,6 +1,6 @@
 import json
 import os
-from typing import Any, Dict
+from typing import Any, Dict, Generator
 
 import click
 
@@ -34,3 +34,16 @@ def json_dumps(data: Any, humand_readable: bool = False) -> str:
         return json.dumps(data, indent=4, sort_keys=True)
     else:
         return json.dumps(data, separators=(",", ":"))
+
+
+def walk(
+    directory: str, recursive: bool = False, match_directories: bool = False
+) -> Generator[os.DirEntry, None, None]:
+    for entry in os.scandir(directory):
+        is_dir = entry.is_dir()
+        if recursive and is_dir:
+            yield from walk(entry.path, recursive, match_directories)
+        elif entry.is_file() or (match_directories and is_dir):
+            yield entry
+        else:
+            continue