seating.seating

  1import argparse
  2
  3import ast_comments as ast
  4import black
  5from pathier import Pathier
  6
  7
  8def get_seat_sections(source: str) -> list[tuple[int, int]]:
  9    """Return a list of line number pairs for content between `# Seat` comments in `source`.
 10
 11    If `source` has no `# Seat` comments, a list with one tuple will be returned: `[(1, number_of_lines_in_source)]`"""
 12
 13    if "# Seat" in source:
 14        lines = source.splitlines()
 15        sections = []
 16        previous_endline = lambda: sections[-1][1]
 17        for i, line in enumerate(lines):
 18            if "# Seat" in line:
 19                if not sections:
 20                    sections = [(1, i + 1)]
 21                else:
 22                    sections.append((previous_endline() + 1, i + 1))
 23        sections.append((previous_endline() + 1, len(lines) + 1))
 24        return sections
 25    return [(1, len(source.splitlines()) + 1)]
 26
 27
 28class Seats:
 29    def __init__(self):
 30        self.before = []
 31        self.assigns = []
 32        self.dunders = []
 33        self.properties = []
 34        self.functions = []
 35        self.after = []
 36        self.seats = []
 37        # These will be a list of tuples containing the node and the index it was found at
 38        # so they can be reinserted after sorting
 39        self.expressions = []
 40        self.comments = []
 41
 42    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
 43        return sorted(nodes, key=lambda node: node.name)
 44
 45    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
 46        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
 47        dunders = self.sort_nodes_by_name(dunders)
 48        init = None
 49        for i, dunder in enumerate(dunders):
 50            if dunder.name == "__init__":
 51                init = dunders.pop(i)
 52                break
 53        if init:
 54            dunders.insert(0, init)
 55        return dunders
 56
 57    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
 58        """Sort assignment statments."""
 59
 60        def get_name(node: ast.stmt) -> str:
 61            type_ = type(node)
 62            if type_ == ast.Assign:
 63                return node.targets[0].id
 64            else:
 65                return node.target.id
 66
 67        return sorted(assigns, key=get_name)
 68
 69    def sort(self) -> list[ast.stmt]:
 70        """Sort and return members as a single list."""
 71        self.dunders = self.sort_dunders(self.dunders)
 72        self.functions = self.sort_nodes_by_name(self.functions)
 73        self.properties = self.sort_nodes_by_name(self.properties)
 74        self.assigns = self.sort_assigns(self.assigns)
 75        body = (
 76            self.before
 77            + self.assigns
 78            + self.dunders
 79            + self.properties
 80            + self.functions
 81            + self.seats
 82            + self.after
 83        )
 84        for expression in self.expressions + self.comments:
 85            body.insert(expression[1], expression[0])
 86        return body
 87
 88
 89def seat(
 90    source: str, start_line: int | None = None, stop_line: int | None = None
 91) -> str:
 92    """Sort the contents of classes in `source`, where `source` is parsable Python code.
 93    Anything not inside a class will be untouched.
 94
 95    The modified `source` will be returned.
 96
 97    #### :params:
 98
 99    * `start_line`: Only sort contents after this line.
100
101    * `stop_line`: Only sort contents before this line.
102
103    If you have class contents that are grouped a certain way and you want the groups individually sorted
104    so that the grouping is maintained, you can use `# Seat` to demarcate the groups.
105
106    i.e. if the source is:
107    >>> class MyClass():
108    >>>     {arbitrary lines of code}
109    >>>     # Seat
110    >>>     {more arbitrary code}
111    >>>     # Seat
112    >>>     {yet more code}
113
114    Then the three sets of code in brackets will be sorted independently from one another
115    (assuming no values are given for `start_line` or `stop_line`).
116
117    #### :Sorting and Priority:
118
119    * Class variables declared in class body outside of a function
120    * Dunder methods
121    * Functions decorated with `property` or corresponding `.setter` and `.deleter` methods
122    * Class functions
123
124    Each of these groups will be sorted alphabetically with respect to themselves.
125
126    The only exception is for dunder methods.
127    They will be sorted alphabetically except that `__init__` will be first.
128    """
129    tree = ast.parse(source, type_comments=True)
130    start_line = start_line or 0
131    stop_line = stop_line or len(source.splitlines()) + 1
132    sections = get_seat_sections(source)
133    for section in sections:
134        for i, stmt in enumerate(tree.body):
135            if type(stmt) == ast.ClassDef:
136                order = Seats()
137                for j, child in enumerate(stmt.body):
138                    try:
139                        type_ = type(child)
140                        if child.lineno <= start_line or child.lineno < section[0]:
141                            order.before.append(child)
142                        elif stop_line < child.lineno or child.lineno > section[1]:
143                            order.after.append(child)
144                        elif type_ == ast.Expr:
145                            order.expressions.append((child, j))
146                        elif type_ == ast.Comment:
147                            if "# Seat" in child.value:
148                                order.seats.append(child)
149                            else:
150                                order.comments.append((child, j))
151                        elif type_ in [ast.Assign, ast.AugAssign, ast.AnnAssign]:
152                            order.assigns.append(child)
153                        elif child.name.startswith("__") and child.name.endswith("__"):
154                            order.dunders.append(child)
155                        elif child.decorator_list:
156                            for decorator in child.decorator_list:
157                                decorator_type = type(decorator)
158                                if (
159                                    decorator_type == ast.Name
160                                    and "property" in decorator.id
161                                ) or (
162                                    decorator_type == ast.Attribute
163                                    and decorator.attr in ["setter", "deleter"]
164                                ):
165                                    order.properties.append(child)
166                                    break
167                            if child not in order.properties:
168                                order.functions.append(child)
169                        else:
170                            order.functions.append(child)
171                    except Exception as e:
172                        print(ast.dump(child, indent=2))
173                        raise e
174                tree.body[i].body = order.sort()
175    return ast.unparse(tree)
176
177
178def get_args() -> argparse.Namespace:
179    parser = argparse.ArgumentParser()
180
181    parser.add_argument("file", type=str, help=""" The file to format. """)
182    parser.add_argument(
183        "--start",
184        type=int,
185        default=None,
186        help=""" Optional line number to start formatting at. """,
187    )
188    parser.add_argument(
189        "--stop",
190        type=int,
191        default=None,
192        help=""" Optional line number to stop formatting at. """,
193    )
194    parser.add_argument(
195        "-nb",
196        "--noblack",
197        action="store_true",
198        help=""" Don't format file with Black after sorting. """,
199    )
200    parser.add_argument(
201        "-o",
202        "--output",
203        default=None,
204        help=""" Write changes to this file, otherwise changes are written back to the original file. """,
205    )
206    parser.add_argument(
207        "-d",
208        "--dump",
209        action="store_true",
210        help=""" Dump ast tree to file instead of doing anything else. 
211        For debugging purposes.""",
212    )
213    args = parser.parse_args()
214
215    return args
216
217
218def main(args: argparse.Namespace | None = None):
219    if not args:
220        args = get_args()
221    source = Pathier(args.file).read_text()
222    if args.dump:
223        file = Pathier(args.file)
224        file = file.with_name(f"{file.stem}_ast_dump.txt").write_text(
225            ast.dump(ast.parse(source, type_comments=True), indent=2)
226        )
227    else:
228        source = seat(source, args.start, args.stop)
229        if not args.noblack:
230            source = black.format_str(source, mode=black.Mode())
231        Pathier(args.output or args.file).write_text(source)
232
233
234if __name__ == "__main__":
235    main(get_args())