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 fix_type_ignore(source: str) -> str:
 90    """Fix misplacement of `# type: ignore` comments in sorted `source`.
 91
 92    Corrects when
 93    >>> var = 6 # type: ignore
 94
 95    gets turned into
 96
 97    >>> # type: ignore
 98    >>> var = 6"""
 99    lines = source.splitlines(True)
100    clean = lambda s: s.replace(" ", "").replace("\n", "")
101    for i, line in enumerate(lines):
102        if clean(line) == "#type:ignore" and 0 < i < len(lines):
103            lines[i] = ""
104            if "#type:ignore" not in clean(lines[i + 1]):
105                lines[i + 1] = lines[i + 1].strip("\n") + "# type: ignore\n"
106    return "".join(lines)
107
108
109def seat(
110    source: str, start_line: int | None = None, stop_line: int | None = None
111) -> str:
112    """Sort the contents of classes in `source`, where `source` is parsable Python code.
113    Anything not inside a class will be untouched.
114
115    The modified `source` will be returned.
116
117    #### :params:
118
119    * `start_line`: Only sort contents after this line.
120
121    * `stop_line`: Only sort contents before this line.
122
123    If you have class contents that are grouped a certain way and you want the groups individually sorted
124    so that the grouping is maintained, you can use `# Seat` to demarcate the groups.
125
126    i.e. if the source is:
127    >>> class MyClass():
128    >>>     {arbitrary lines of code}
129    >>>     # Seat
130    >>>     {more arbitrary code}
131    >>>     # Seat
132    >>>     {yet more code}
133
134    Then the three sets of code in brackets will be sorted independently from one another
135    (assuming no values are given for `start_line` or `stop_line`).
136
137    #### :Sorting and Priority:
138
139    * Class variables declared in class body outside of a function
140    * Dunder methods
141    * Functions decorated with `property` or corresponding `.setter` and `.deleter` methods
142    * Class functions
143
144    Each of these groups will be sorted alphabetically with respect to themselves.
145
146    The only exception is for dunder methods.
147    They will be sorted alphabetically except that `__init__` will be first.
148    """
149    tree = ast.parse(source, type_comments=True)
150    start_line = start_line or 0
151    stop_line = stop_line or len(source.splitlines()) + 1
152    sections = get_seat_sections(source)
153    for section in sections:
154        for i, stmt in enumerate(tree.body):
155            if type(stmt) == ast.ClassDef:
156                order = Seats()
157                for j, child in enumerate(stmt.body):
158                    try:
159                        type_ = type(child)
160                        if child.lineno <= start_line or child.lineno < section[0]:
161                            order.before.append(child)
162                        elif stop_line < child.lineno or child.lineno > section[1]:
163                            order.after.append(child)
164                        elif type_ == ast.Expr:
165                            order.expressions.append((child, j))
166                        elif type_ == ast.Comment:
167                            if "# Seat" in child.value:
168                                order.seats.append(child)
169                            else:
170                                order.comments.append((child, j))
171                        elif type_ in [ast.Assign, ast.AugAssign, ast.AnnAssign]:
172                            order.assigns.append(child)
173                        elif child.name.startswith("__") and child.name.endswith("__"):
174                            order.dunders.append(child)
175                        elif child.decorator_list:
176                            for decorator in child.decorator_list:
177                                decorator_type = type(decorator)
178                                if (
179                                    decorator_type == ast.Name
180                                    and "property" in decorator.id
181                                ) or (
182                                    decorator_type == ast.Attribute
183                                    and decorator.attr in ["setter", "deleter"]
184                                ):
185                                    order.properties.append(child)
186                                    break
187                            if child not in order.properties:
188                                order.functions.append(child)
189                        else:
190                            order.functions.append(child)
191                    except Exception as e:
192                        print(ast.dump(child, indent=2))
193                        raise e
194                tree.body[i].body = order.sort()
195    source = ast.unparse(tree)
196    return fix_type_ignore(source)
197
198
199def get_args() -> argparse.Namespace:
200    parser = argparse.ArgumentParser()
201
202    parser.add_argument("file", type=str, help=""" The file to format. """)
203    parser.add_argument(
204        "--start",
205        type=int,
206        default=None,
207        help=""" Optional line number to start formatting at. """,
208    )
209    parser.add_argument(
210        "--stop",
211        type=int,
212        default=None,
213        help=""" Optional line number to stop formatting at. """,
214    )
215    parser.add_argument(
216        "-nb",
217        "--noblack",
218        action="store_true",
219        help=""" Don't format file with Black after sorting. """,
220    )
221    parser.add_argument(
222        "-o",
223        "--output",
224        default=None,
225        help=""" Write changes to this file, otherwise changes are written back to the original file. """,
226    )
227    parser.add_argument(
228        "-d",
229        "--dump",
230        action="store_true",
231        help=""" Dump ast tree to file instead of doing anything else. 
232        For debugging purposes.""",
233    )
234    args = parser.parse_args()
235
236    return args
237
238
239def main(args: argparse.Namespace | None = None):
240    if not args:
241        args = get_args()
242    source = Pathier(args.file).read_text()
243    if args.dump:
244        file = Pathier(args.file)
245        file = file.with_name(f"{file.stem}_ast_dump.txt").write_text(
246            ast.dump(ast.parse(source, type_comments=True), indent=2)
247        )
248    else:
249        source = seat(source, args.start, args.stop)
250        if not args.noblack:
251            source = black.format_str(source, mode=black.Mode())
252        Pathier(args.output or args.file).write_text(source)
253
254
255if __name__ == "__main__":
256    main(get_args())
def get_seat_sections(source: str) -> list[tuple[int, int]]:
 9def get_seat_sections(source: str) -> list[tuple[int, int]]:
10    """Return a list of line number pairs for content between `# Seat` comments in `source`.
11
12    If `source` has no `# Seat` comments, a list with one tuple will be returned: `[(1, number_of_lines_in_source)]`"""
13
14    if "# Seat" in source:
15        lines = source.splitlines()
16        sections = []
17        previous_endline = lambda: sections[-1][1]
18        for i, line in enumerate(lines):
19            if "# Seat" in line:
20                if not sections:
21                    sections = [(1, i + 1)]
22                else:
23                    sections.append((previous_endline() + 1, i + 1))
24        sections.append((previous_endline() + 1, len(lines) + 1))
25        return sections
26    return [(1, len(source.splitlines()) + 1)]

Return a list of line number pairs for content between # Seat comments in source.

If source has no # Seat comments, a list with one tuple will be returned: [(1, number_of_lines_in_source)]

class Seats:
29class Seats:
30    def __init__(self):
31        self.before = []
32        self.assigns = []
33        self.dunders = []
34        self.properties = []
35        self.functions = []
36        self.after = []
37        self.seats = []
38        # These will be a list of tuples containing the node and the index it was found at
39        # so they can be reinserted after sorting
40        self.expressions = []
41        self.comments = []
42
43    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
44        return sorted(nodes, key=lambda node: node.name)
45
46    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
47        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
48        dunders = self.sort_nodes_by_name(dunders)
49        init = None
50        for i, dunder in enumerate(dunders):
51            if dunder.name == "__init__":
52                init = dunders.pop(i)
53                break
54        if init:
55            dunders.insert(0, init)
56        return dunders
57
58    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
59        """Sort assignment statments."""
60
61        def get_name(node: ast.stmt) -> str:
62            type_ = type(node)
63            if type_ == ast.Assign:
64                return node.targets[0].id
65            else:
66                return node.target.id
67
68        return sorted(assigns, key=get_name)
69
70    def sort(self) -> list[ast.stmt]:
71        """Sort and return members as a single list."""
72        self.dunders = self.sort_dunders(self.dunders)
73        self.functions = self.sort_nodes_by_name(self.functions)
74        self.properties = self.sort_nodes_by_name(self.properties)
75        self.assigns = self.sort_assigns(self.assigns)
76        body = (
77            self.before
78            + self.assigns
79            + self.dunders
80            + self.properties
81            + self.functions
82            + self.seats
83            + self.after
84        )
85        for expression in self.expressions + self.comments:
86            body.insert(expression[1], expression[0])
87        return body
Seats()
30    def __init__(self):
31        self.before = []
32        self.assigns = []
33        self.dunders = []
34        self.properties = []
35        self.functions = []
36        self.after = []
37        self.seats = []
38        # These will be a list of tuples containing the node and the index it was found at
39        # so they can be reinserted after sorting
40        self.expressions = []
41        self.comments = []
def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
43    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
44        return sorted(nodes, key=lambda node: node.name)
def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
46    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
47        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
48        dunders = self.sort_nodes_by_name(dunders)
49        init = None
50        for i, dunder in enumerate(dunders):
51            if dunder.name == "__init__":
52                init = dunders.pop(i)
53                break
54        if init:
55            dunders.insert(0, init)
56        return dunders

Sort dunders alphabetically, except __init__ is placed at the front, if it exists.

def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
58    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
59        """Sort assignment statments."""
60
61        def get_name(node: ast.stmt) -> str:
62            type_ = type(node)
63            if type_ == ast.Assign:
64                return node.targets[0].id
65            else:
66                return node.target.id
67
68        return sorted(assigns, key=get_name)

Sort assignment statments.

def sort(self) -> list[ast.stmt]:
70    def sort(self) -> list[ast.stmt]:
71        """Sort and return members as a single list."""
72        self.dunders = self.sort_dunders(self.dunders)
73        self.functions = self.sort_nodes_by_name(self.functions)
74        self.properties = self.sort_nodes_by_name(self.properties)
75        self.assigns = self.sort_assigns(self.assigns)
76        body = (
77            self.before
78            + self.assigns
79            + self.dunders
80            + self.properties
81            + self.functions
82            + self.seats
83            + self.after
84        )
85        for expression in self.expressions + self.comments:
86            body.insert(expression[1], expression[0])
87        return body

Sort and return members as a single list.

def fix_type_ignore(source: str) -> str:
 90def fix_type_ignore(source: str) -> str:
 91    """Fix misplacement of `# type: ignore` comments in sorted `source`.
 92
 93    Corrects when
 94    >>> var = 6 # type: ignore
 95
 96    gets turned into
 97
 98    >>> # type: ignore
 99    >>> var = 6"""
100    lines = source.splitlines(True)
101    clean = lambda s: s.replace(" ", "").replace("\n", "")
102    for i, line in enumerate(lines):
103        if clean(line) == "#type:ignore" and 0 < i < len(lines):
104            lines[i] = ""
105            if "#type:ignore" not in clean(lines[i + 1]):
106                lines[i + 1] = lines[i + 1].strip("\n") + "# type: ignore\n"
107    return "".join(lines)

Fix misplacement of # type: ignore comments in sorted source.

Corrects when

>>> var = 6 # type: ignore

gets turned into

>>> # type: ignore
>>> var = 6
def seat( source: str, start_line: int | None = None, stop_line: int | None = None) -> str:
110def seat(
111    source: str, start_line: int | None = None, stop_line: int | None = None
112) -> str:
113    """Sort the contents of classes in `source`, where `source` is parsable Python code.
114    Anything not inside a class will be untouched.
115
116    The modified `source` will be returned.
117
118    #### :params:
119
120    * `start_line`: Only sort contents after this line.
121
122    * `stop_line`: Only sort contents before this line.
123
124    If you have class contents that are grouped a certain way and you want the groups individually sorted
125    so that the grouping is maintained, you can use `# Seat` to demarcate the groups.
126
127    i.e. if the source is:
128    >>> class MyClass():
129    >>>     {arbitrary lines of code}
130    >>>     # Seat
131    >>>     {more arbitrary code}
132    >>>     # Seat
133    >>>     {yet more code}
134
135    Then the three sets of code in brackets will be sorted independently from one another
136    (assuming no values are given for `start_line` or `stop_line`).
137
138    #### :Sorting and Priority:
139
140    * Class variables declared in class body outside of a function
141    * Dunder methods
142    * Functions decorated with `property` or corresponding `.setter` and `.deleter` methods
143    * Class functions
144
145    Each of these groups will be sorted alphabetically with respect to themselves.
146
147    The only exception is for dunder methods.
148    They will be sorted alphabetically except that `__init__` will be first.
149    """
150    tree = ast.parse(source, type_comments=True)
151    start_line = start_line or 0
152    stop_line = stop_line or len(source.splitlines()) + 1
153    sections = get_seat_sections(source)
154    for section in sections:
155        for i, stmt in enumerate(tree.body):
156            if type(stmt) == ast.ClassDef:
157                order = Seats()
158                for j, child in enumerate(stmt.body):
159                    try:
160                        type_ = type(child)
161                        if child.lineno <= start_line or child.lineno < section[0]:
162                            order.before.append(child)
163                        elif stop_line < child.lineno or child.lineno > section[1]:
164                            order.after.append(child)
165                        elif type_ == ast.Expr:
166                            order.expressions.append((child, j))
167                        elif type_ == ast.Comment:
168                            if "# Seat" in child.value:
169                                order.seats.append(child)
170                            else:
171                                order.comments.append((child, j))
172                        elif type_ in [ast.Assign, ast.AugAssign, ast.AnnAssign]:
173                            order.assigns.append(child)
174                        elif child.name.startswith("__") and child.name.endswith("__"):
175                            order.dunders.append(child)
176                        elif child.decorator_list:
177                            for decorator in child.decorator_list:
178                                decorator_type = type(decorator)
179                                if (
180                                    decorator_type == ast.Name
181                                    and "property" in decorator.id
182                                ) or (
183                                    decorator_type == ast.Attribute
184                                    and decorator.attr in ["setter", "deleter"]
185                                ):
186                                    order.properties.append(child)
187                                    break
188                            if child not in order.properties:
189                                order.functions.append(child)
190                        else:
191                            order.functions.append(child)
192                    except Exception as e:
193                        print(ast.dump(child, indent=2))
194                        raise e
195                tree.body[i].body = order.sort()
196    source = ast.unparse(tree)
197    return fix_type_ignore(source)

Sort the contents of classes in source, where source is parsable Python code. Anything not inside a class will be untouched.

The modified source will be returned.

:params:

  • start_line: Only sort contents after this line.

  • stop_line: Only sort contents before this line.

If you have class contents that are grouped a certain way and you want the groups individually sorted so that the grouping is maintained, you can use # Seat to demarcate the groups.

i.e. if the source is:

>>> class MyClass():
>>>     {arbitrary lines of code}
>>>     # Seat
>>>     {more arbitrary code}
>>>     # Seat
>>>     {yet more code}

Then the three sets of code in brackets will be sorted independently from one another (assuming no values are given for start_line or stop_line).

:Sorting and Priority:

  • Class variables declared in class body outside of a function
  • Dunder methods
  • Functions decorated with property or corresponding .setter and .deleter methods
  • Class functions

Each of these groups will be sorted alphabetically with respect to themselves.

The only exception is for dunder methods. They will be sorted alphabetically except that __init__ will be first.

def get_args() -> argparse.Namespace:
200def get_args() -> argparse.Namespace:
201    parser = argparse.ArgumentParser()
202
203    parser.add_argument("file", type=str, help=""" The file to format. """)
204    parser.add_argument(
205        "--start",
206        type=int,
207        default=None,
208        help=""" Optional line number to start formatting at. """,
209    )
210    parser.add_argument(
211        "--stop",
212        type=int,
213        default=None,
214        help=""" Optional line number to stop formatting at. """,
215    )
216    parser.add_argument(
217        "-nb",
218        "--noblack",
219        action="store_true",
220        help=""" Don't format file with Black after sorting. """,
221    )
222    parser.add_argument(
223        "-o",
224        "--output",
225        default=None,
226        help=""" Write changes to this file, otherwise changes are written back to the original file. """,
227    )
228    parser.add_argument(
229        "-d",
230        "--dump",
231        action="store_true",
232        help=""" Dump ast tree to file instead of doing anything else. 
233        For debugging purposes.""",
234    )
235    args = parser.parse_args()
236
237    return args
def main(args: argparse.Namespace | None = None):
240def main(args: argparse.Namespace | None = None):
241    if not args:
242        args = get_args()
243    source = Pathier(args.file).read_text()
244    if args.dump:
245        file = Pathier(args.file)
246        file = file.with_name(f"{file.stem}_ast_dump.txt").write_text(
247            ast.dump(ast.parse(source, type_comments=True), indent=2)
248        )
249    else:
250        source = seat(source, args.start, args.stop)
251        if not args.noblack:
252            source = black.format_str(source, mode=black.Mode())
253        Pathier(args.output or args.file).write_text(source)