dj_model_strip.py

import ast
import astor
import tkinter as tk
from tkinter import scrolledtext, IntVar, Checkbutton

class EllipsisTransformer(ast.NodeTransformer):
    def __init__(self, remove_args, remove_methods):
        self.remove_args = remove_args
        self.remove_methods = remove_methods

    def visit_FunctionDef(self, node):
        if self.remove_methods:
            return None  # removes the node
        else:
            self.replace_body_with_ellipsis(node)
            return node

    def visit_AsyncFunctionDef(self, node):
        if self.remove_methods:
            return None  # removes the node
        else:
            self.replace_body_with_ellipsis(node)
            return node

    def visit_ClassDef(self, node):
        if not self.remove_methods:
            decorators_to_keep = []
            for decorator in node.decorator_list:
                if isinstance(decorator, ast.Name) and decorator.id in ["property", "cached_property"]:
                    decorators_to_keep.append(decorator)
            node.decorator_list = decorators_to_keep
        self.generic_visit(node)
        return node

    def visit_Import(self, node):
        return None  # removes the node

    def visit_ImportFrom(self, node):
        return None  # removes the node

    def visit_Call(self, node):
        self.generic_visit(node)
        if self.remove_args:
            if isinstance(node.func, ast.Attribute):
                if isinstance(node.func.value, ast.Name) and node.func.value.id == 'models':
                    node.args = [ast.Constant(value="...")]
                    node.keywords = []
        return node

    def replace_body_with_ellipsis(self, node):
        for child in ast.iter_child_nodes(node):
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                continue
        node.body = [ast.Expr(value=ast.Ellipsis())]

class Application(tk.Frame):
    def __init__(self, master=None):
        super().__init__(master)
        self.master = master
        self.grid()
        self.create_widgets()

    def create_widgets(self):
        self.input_text = scrolledtext.ScrolledText(self)
        self.input_text.grid(row=0, column=0, columnspan=2)

        self.summarize_button = tk.Button(self)
        self.summarize_button["text"] = "Summarize"
        self.summarize_button["command"] = self.summarize
        self.summarize_button.grid(row=1, column=0)

        self.quit_button = tk.Button(self, text="QUIT", command=self.master.destroy)
        self.quit_button.grid(row=1, column=1)

        self.remove_args = IntVar()
        self.remove_args_check = Checkbutton(self, text="Remove field arguments", variable=self.remove_args)
        self.remove_args_check.grid(row=2, column=0, columnspan=2)

        self.remove_methods = IntVar()
        self.remove_methods_check = Checkbutton(self, text="Remove methods", variable=self.remove_methods)
        self.remove_methods_check.grid(row=3, column=0, columnspan=2)

    def summarize(self):
        source = self.input_text.get("1.0", "end")
        module = ast.parse(source)
        transformer = EllipsisTransformer(self.remove_args.get(), self.remove_methods.get())
        transformer.visit(module)
        summarized_code = astor.to_source(module)
        self.input_text.delete('1.0', "end")
        self.input_text.insert("1.0", summarized_code)

root = tk.Tk()
app = Application(master=root)
app.mainloop()