Add variable fitler to snapshot

CHANGELOG.md

@@ -10,6 +10,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 ### ✨ Enhancements
 
 - Added `include` filter to `snapshot`
+- Added `exclude_vars` filter to `snapshot`
 
 ### 💫 New checkers
 

python_ta/debug/snapshot.py

@@ -37,11 +37,25 @@ def get_filtered_global_variables(frame: FrameType) -> dict:
     return {"__main__": true_global_vars}
 
 
+def get_filtered_local_variables(frame, exclude_vars: Optional[Iterable[str | re.Pattern]]) -> dict:
+    """
+    Helper function for filtering local variables in a frame.
+    """
+    if exclude_vars:
+        return {
+            var: frame.f_locals[var]
+            for var in frame.f_locals
+            if not any(re.search(regex, var) for regex in exclude_vars)
+        }
+    return frame.f_locals
+
+
 def snapshot(
     save: bool = False,
     memory_viz_args: Optional[list[str]] = None,
     memory_viz_version: str = "latest",
     include: Optional[Iterable[str | re.Pattern]] = None,
+    exclude_vars: Optional[Iterable[str | re.Pattern]] = None,
 ):
     """Capture a snapshot of local variables from the current and outer stack frames
     where the 'snapshot' function is called. Returns a list of dictionaries,
@@ -56,14 +70,17 @@ def snapshot(
     include can be used to specify a collection of function names, either as strings or regular expressions,
     whose variables will be captured. By default, all variables in all functions will be captured if no `include`
     argument is provided.
+    exclude_vars can be used to specify a collection of variable names, either as strings or regular expressions,
+    that will be excluded from the snapshot. By default, all variables will be captured if no `exclude_vars` is provided.
     """
     variables = []
     frame = inspect.currentframe().f_back
 
     while frame:
         if include is None or any(re.search(regex, frame.f_code.co_name) for regex in include):
             if frame.f_code.co_name != "<module>":
-                variables.append({frame.f_code.co_name: frame.f_locals})
+                local_vars = get_filtered_local_variables(frame, exclude_vars)
+                variables.append({frame.f_code.co_name: local_vars})
             else:
                 global_vars = get_filtered_global_variables(frame)
                 variables.append(global_vars)

tests/test_debug/test_snapshot.py

@@ -54,6 +54,59 @@ def func_cyclic() -> list:
     return snapshot()
 
 
+def func_with_include(include: Optional[Iterable[str | re.Pattern]] = None) -> list[dict]:
+    """
+    Function for snapshot() testing with include parameter.
+    """
+    test_var1a = "David is cool!"
+    test_var2a = "Students Developing Software"
+    return snapshot(include=include)
+
+
+def func_with_include_nested(include: Optional[Iterable[str | re.Pattern]] = None) -> list[dict]:
+    """
+    Function for snapshot() testing with include parameter.
+    """
+    test_var1b = {"SDS_coolest_project": "PyTA"}
+    test_var2b = ("Leo", "tester")
+    return func_with_include(include=include)
+
+
+def func_with_unserializable_objects() -> list[dict]:
+    """
+    Function for snapshot() testing with unserializable objects.
+    """
+    path = pathlib.PosixPath("some path")
+    vars_in_curr_func = [snapshot()[0]]
+    processed_result = snapshot_to_json(vars_in_curr_func)
+    json.dumps(processed_result)
+    return processed_result
+
+
+def func_with_exclude(
+    exclude_vars: Optional[Iterable[str | re.Pattern]] = None,
+    include: Optional[Iterable[str | re.Pattern]] = None,
+) -> list[dict]:
+    """
+    Function for snapshot() testing with exclude parameter.
+    """
+    test_var1a = "David is cool!"
+    test_var2a = "Students Developing Software"
+    return snapshot(exclude_vars=exclude_vars, include=include)
+
+
+def func_with_exclude_nested(
+    exclude_vars: Optional[Iterable[str | re.Pattern]] = None,
+    include: Optional[Iterable[str | re.Pattern]] = None,
+) -> list[dict]:
+    """
+    Function for snapshot() testing with exclude parameter.
+    """
+    test_var1b = {"SDS_coolest_project": "PyTA"}
+    test_var2b = ("Leo", "tester")
+    return func_with_exclude(exclude_vars=exclude_vars, include=include)
+
+
 def test_snapshot_one_level() -> None:
     """
     Examines whether the snapshot() function accurately captures
@@ -650,26 +703,6 @@ def test_snapshot_save_stdout():
     assert result.stdout == expected_svg
 
 
-def func_with_include(include: Optional[Iterable[str | re.Pattern]] = None) -> list[dict]:
-    test_var1a = "David is cool!"
-    test_var2a = "Students Developing Software"
-    return snapshot(include=include)
-
-
-def func_with_include_nested(include: Optional[Iterable[str | re.Pattern]] = None) -> list[dict]:
-    test_var1b = {"SDS_coolest_project": "PyTA"}
-    test_var2b = ("Leo", "tester")
-    return func_with_include(include=include)
-
-
-def func_with_unserializable_objects() -> list[dict]:
-    path = pathlib.PosixPath("some path")
-    vars_in_curr_func = [snapshot()[0]]
-    processed_result = snapshot_to_json(vars_in_curr_func)
-    json.dumps(processed_result)
-    return processed_result
-
-
 def test_snapshot_only_includes_function_self():
     result = func_with_include(include=("func_with_include",))
     assert result == [
@@ -738,3 +771,88 @@ def test_snapshot_serializes_unserializable_value():
         },
         {"id": 1, "type": "PosixPath", "value": repr(pathlib.PosixPath("some path"))},
     ]
+
+
+def test_snapshot_excludes_one_variable_in_current_frame():
+    """
+    Test snapshot() excludes one variable in the current frame.
+    """
+    result = func_with_exclude(exclude_vars=["test_var1a"], include=["func_with_exclude"])
+    assert result == [
+        {
+            "func_with_exclude": {
+                "exclude_vars": ["test_var1a"],
+                "include": ["func_with_exclude"],
+                "test_var2a": "Students Developing Software",
+            }
+        }
+    ]
+
+
+def test_snapshot_excludes_multiple_variables():
+    """
+    Test snapshot() excludes multiple variables in the current frame.
+    """
+    result = func_with_exclude(
+        exclude_vars=["test_var1a", "test_var2a"], include=["func_with_exclude"]
+    )
+    assert result == [
+        {
+            "func_with_exclude": {
+                "exclude_vars": ["test_var1a", "test_var2a"],
+                "include": ["func_with_exclude"],
+            }
+        }
+    ]
+
+
+def test_snapshot_excludes_variables_in_nested_frames():
+    """
+    Test snapshot() excludes variables in nested frames.
+    """
+    result = func_with_exclude_nested(
+        exclude_vars=["test_var1b"], include=["func_with_exclude", "func_with_exclude_nested"]
+    )
+    assert result == [
+        {
+            "func_with_exclude": {
+                "exclude_vars": ["test_var1b"],
+                "include": ["func_with_exclude", "func_with_exclude_nested"],
+                "test_var1a": "David is cool!",
+                "test_var2a": "Students Developing Software",
+            }
+        },
+        {
+            "func_with_exclude_nested": {
+                "exclude_vars": ["test_var1b"],
+                "include": ["func_with_exclude", "func_with_exclude_nested"],
+                "test_var2b": ("Leo", "tester"),
+            }
+        },
+    ]
+
+
+def test_snapshot_excludes_variables_with_regex():
+    """
+    Test snapshot() excludes variables in nested frames using regex.
+    """
+    result = func_with_exclude_nested(
+        exclude_vars=[re.compile("test_var1.*")],
+        include=["func_with_exclude", "func_with_exclude_nested"],
+    )
+    assert result == [
+        {
+            "func_with_exclude": {
+                "exclude_vars": [re.compile("test_var1.*")],
+                "include": ["func_with_exclude", "func_with_exclude_nested"],
+                "test_var2a": "Students Developing Software",
+            }
+        },
+        {
+            "func_with_exclude_nested": {
+                "exclude_vars": [re.compile("test_var1.*")],
+                "include": ["func_with_exclude", "func_with_exclude_nested"],
+                "test_var2b": ("Leo", "tester"),
+            }
+        },
+    ]