feat: filter_code

README.md

@@ -2,18 +2,16 @@
 
 A place to dump might-be-useful-again code as an alternative of leaving in a notebook where it will never be found again
 
-To install:	```pip install titbit```
-
+To install: `pip install titbit`
 
 # Examples
 
-
 ## git_action_on_projects
 
 Take git actions all the projects in the list of projects.
 A project can be a folder path or a module, or the name of a module/package.
 
-Tip: Use `functools.partial` to set the `action`, `on_error` and `egress` and get 
+Tip: Use `functools.partial` to set the `action`, `on_error` and `egress` and get
 the function you need to perform bulk actions.
 
 Usage:
@@ -30,7 +28,7 @@ Usage:
 ## mermaid_to_graphviz
 
 Converts mermaid code to graphviz code.
-    
+
 ```python
 >>> from titbit import mermaid_to_graphviz
 >>> mermaid_code = '''
@@ -48,10 +46,36 @@ graph TD
 <BLANKLINE>
 }
 ```
+
+## filter_code
+
+Produce a version of the code that is filtered by the filt filter function.
+
+```python
+>>> from titbit import filter_code
+>>> code_str = '''
+... a = b + c
+... print(a)
+... d[0] = a
+... for i in range(1, 10):
+...     d[i] = d[i-1] + b
+... def f(x):
+...     return x + 1
+... e = [d, a, f]
+... '''
+>>> import ast
+>>> filtered_code_str = filter_code(code_str, lambda x: isinstance(x, ast.Assign))
+>>> assert filtered_code_str.strip() == ('''
+... a = b + c
+... d[0] = a
+... e = [d, a, f]
+... '''.strip())
+```
+
 ## bound_properties_refactor
 
 Generate code that refactors "flat code" into a reusable "controller" class.
-Also checkout the `BoundPropertiesRefactor` class that does all the work: 
+Also checkout the `BoundPropertiesRefactor` class that does all the work:
 With it, you'll be able to compute intermediate datas that may be of interest.
 
 ```python
@@ -89,4 +113,4 @@ def ensure_ast(code: AST) -> AST:
     If input is a string, parses it as Python code and returns the resulting AST.
     If the input is a module object, it will get the code, parse it, and return an AST.
     """
-```
+```

titbit/__init__.py

@@ -6,9 +6,14 @@
 """
 
 from titbit.ast_utils import (
-    rename_variables, bound_properties_refactor, BoundPropertiesRefactor, ensure_ast
+    rename_variables,
+    bound_properties_refactor,
+    BoundPropertiesRefactor,
+    ensure_ast,
+    filter_code,
 )
 
+
 def git_action_on_projects(
     projects,
     action='pull',
@@ -75,7 +80,7 @@ def mermaid_to_graphviz(
     mermaid_code, extra_replacements=(), *, prefix='', suffix='', egress=None
 ):
     """Converts mermaid code to graphviz code.
-    
+
     >>> mermaid_code = '''
     ... graph TD
     ... A --> B & C

titbit/ast_utils.py

@@ -70,7 +70,7 @@ def ensure_ast(code: AST) -> AST:
 
     If input is a string, parses it as Python code and returns the resulting AST.
     If the input is a module object, it will get the code, parse it, and return an AST.
-    
+
     """
     assert isinstance(code, AST), "Input must be an AST node or a string."
     return code
@@ -88,6 +88,53 @@ def _(code: ModuleType) -> AST:
     return ensure_ast(inspect.getsource(code))
 
 
+def filter_code(code: Code, filt: Callable[[AST], bool]) -> Code:
+    """
+    Produce a version of the code that is filtered by the filt filter function.
+
+    :param code: The code string or ast.AST object to filter
+    :param filt: A function that takes an AST node and returns True or False
+
+    >>> code_str = '''
+    ... a = b + c
+    ... print(a)
+    ... d[0] = a
+    ... for i in range(1, 10):
+    ...     d[i] = d[i-1] + b
+    ... def f(x):
+    ...     return x + 1
+    ... e = [d, a, f]
+    ... '''
+    >>> import ast
+    >>> filtered_code_str = filter_code(code_str, lambda x: isinstance(x, ast.Assign))
+    >>> assert filtered_code_str.strip() == ('''
+    ... a = b + c
+    ... d[0] = a
+    ... e = [d, a, f]
+    ... '''.strip())
+
+    """
+    input_type = type(code)
+
+    code = ensure_ast(code)
+
+    # walk through the first layer of nodes in the AST and yield the ones that
+    # pass the filter. Do not go deeper than the first layer.
+
+    def _filter_lines(node):
+        for line in node.body:
+            if filt(line):
+                yield line
+
+    # create a new code object with the filtered lines
+    filtered_code = ast.Module(body=list(_filter_lines(code)), type_ignores=[])
+
+    if input_type is str:
+        return ast.unparse(filtered_code)
+
+    return filtered_code
+
+
 def ast_flat_hierarchy():
     """
     Returns a dictionary with the AST classes grouped by their module path.
@@ -270,7 +317,7 @@ class BoundPropertiesRefactor:
     """
     Generate code that refactors "flat code" into a reusable "controller" class.
 
-    You'd usually just use the `bound_properties_refactor` function for this, but 
+    You'd usually just use the `bound_properties_refactor` function for this, but
     this class's instances let's you get intermediate objects that can be useful.
 
     >>> code_str = '''