From 3887640fbdc28a0d4581dad8e6edd12b42363c0a Mon Sep 17 00:00:00 2001
From: Tom de Geus <tom@geus.me>
Date: Wed, 14 Dec 2022 13:20:59 +0100
Subject: [PATCH] shelephant_dump: adding filter on extension

---
 shelephant/cli/shelephant_dump.py | 12 ++++++++++++
 tests/test_main.py                | 12 ++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/shelephant/cli/shelephant_dump.py b/shelephant/cli/shelephant_dump.py
index bff220a..3d971c8 100644
--- a/shelephant/cli/shelephant_dump.py
+++ b/shelephant/cli/shelephant_dump.py
@@ -1,5 +1,6 @@
 import argparse
 import os
+import pathlib
 import re
 import subprocess
 import sys
@@ -23,6 +24,14 @@ def _shelephant_dump_parser():
     parser.add_argument(
         "-e", "--exclude", type=str, action="append", help="Exclude files matching this pattern."
     )
+    parser.add_argument(
+        "-E",
+        "--exclude-extension",
+        type=str,
+        action="append",
+        default=[],
+        help='Exclude files based on their extension (e.g. ".bak").',
+    )
     parser.add_argument("--fmt", type=str, help='Formatter of each line, e.g. ``"mycmd {}"``.')
     parser.add_argument(
         "-c",
@@ -82,6 +91,9 @@ def shelephant_dump(args: list[str]):
             excl = np.logical_or(excl, np.array([re.match(pattern, file) for file in files]))
         files = [file for file, ex in zip(files, excl) if not ex]
 
+    if args.exclude_extension:
+        files = [file for file in files if pathlib.Path(file).suffix not in args.exclude_extension]
+
     if args.sort:
         files = sorted(files)
 
diff --git a/tests/test_main.py b/tests/test_main.py
index 1c837b8..caa8acb 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -6,6 +6,7 @@
 import numpy as np
 
 import shelephant
+from shelephant.cli.shelephant_dump import shelephant_dump
 
 
 def run(cmd, verbose=False):
@@ -240,6 +241,17 @@ def test_append(self):
         os.remove("bar.pdf")
         os.remove("shelephant_dump.yaml")
 
+    def test_exclude(self):
+
+        shelephant_dump(["a.txt", "b.bak", "c.h5", "-E", ".bak", "-o", "dump.yaml"])
+
+        self.assertEqual(
+            shelephant.yaml.read("dump.yaml"),
+            ["a.txt", "c.h5"],
+        )
+
+        os.remove("dump.yaml")
+
 
 class Test_extract(unittest.TestCase):
     def test_single_path(self):