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 Order:
 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
 38    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
 39        return sorted(nodes, key=lambda node: node.name)
 40
 41    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
 42        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
 43        dunders = self.sort_nodes_by_name(dunders)
 44        init = None
 45        for i, dunder in enumerate(dunders):
 46            if dunder.name == "__init__":
 47                init = dunders.pop(i)
 48                break
 49        if init:
 50            dunders.insert(0, init)
 51        return dunders
 52
 53    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
 54        """Sort assignment statments."""
 55
 56        def get_name(node: ast.stmt) -> str:
 57            type_ = type(node)
 58            if type_ == ast.Assign:
 59                return node.targets[0].id
 60            else:
 61                return node.target.id
 62
 63        return sorted(assigns, key=get_name)
 64
 65    def sort(self) -> list[ast.stmt]:
 66        """Sort and return members as a single list."""
 67        self.dunders = self.sort_dunders(self.dunders)
 68        self.functions = self.sort_nodes_by_name(self.functions)
 69        self.properties = self.sort_nodes_by_name(self.properties)
 70        self.assigns = self.sort_assigns(self.assigns)
 71        return (
 72            self.before
 73            + self.assigns
 74            + self.dunders
 75            + self.properties
 76            + self.functions
 77            + self.seats
 78            + self.after
 79        )
 80
 81
 82def seat(
 83    source: str, start_line: int | None = None, stop_line: int | None = None
 84) -> str:
 85    """Sort the contents of classes in `source`, where `source` is parsable Python code.
 86    Anything not inside a class will be untouched.
 87
 88    The modified `source` will be returned.
 89
 90    #### :params:
 91
 92    * `start_line`: Only sort contents after this line.
 93
 94    * `stop_line`: Only sort contents before this line.
 95
 96    If you have class contents that are grouped a certain way and you want the groups individually sorted
 97    so that the grouping is maintained, you can use `# Seat` to demarcate the groups.
 98
 99    Note: Comments that are in a class body, but not in a function will remain at the same line.
100
101    i.e. if the source is:
102    >>> class MyClass():
103    >>>     {arbitrary lines of code}
104    >>>     # Seat
105    >>>     {more arbitrary code}
106    >>>     # Seat
107    >>>     {yet more code}
108
109    Then the three sets of code in brackets will be sorted independently from one another
110    (assuming no values are given for `start_line` or `stop_line`).
111
112    #### :Sorting and Priority:
113
114    * Class variables declared in class body outside of a function
115    * Dunder methods
116    * Functions decorated with `property` or corresponding `.setter` and `.deleter` methods
117    * Class functions
118
119    Each of these groups will be sorted alphabetically with respect to themselves.
120
121    The only exception is for dunder methods.
122    They will be sorted alphabetically except that `__init__` will be first.
123    """
124    tree = ast.parse(source, type_comments=True)
125    start_line = start_line or 0
126    stop_line = stop_line or len(source.splitlines()) + 1
127    sections = get_seat_sections(source)
128    comments = []
129    for section in sections:
130        for i, stmt in enumerate(tree.body):
131            if type(stmt) == ast.ClassDef:
132                order = Order()
133                for child in stmt.body:
134                    type_ = type(child)
135                    if child.lineno <= start_line or child.lineno < section[0]:
136                        order.before.append(child)
137                    elif stop_line < child.lineno or child.lineno > section[1]:
138                        order.after.append(child)
139                    elif type_ == ast.Comment:
140                        if "# Seat" in child.value:
141                            order.seats.append(child)
142                        else:
143                            comments.append(child)
144                    elif type_ in [ast.Assign, ast.AugAssign, ast.AnnAssign]:
145                        order.assigns.append(child)
146                    elif child.name.startswith("__") and child.name.endswith("__"):
147                        order.dunders.append(child)
148                    elif child.decorator_list:
149                        for decorator in child.decorator_list:
150                            decorator_type = type(decorator)
151                            if (
152                                decorator_type == ast.Name
153                                and decorator.id == "property"
154                            ) or (
155                                decorator_type == ast.Attribute
156                                and decorator.attr in ["setter", "deleter"]
157                            ):
158                                order.properties.append(child)
159                                break
160                        if child not in order.properties:
161                            order.functions.append(child)
162                    else:
163                        order.functions.append(child)
164                tree.body[i].body = order.sort()
165    # Put comments back in
166    source = ast.unparse(tree).splitlines()
167    for comment in comments:
168        source.insert(comment.lineno - 1, f"{' '*comment.col_offset}{comment.value}")
169    return "\n".join(source)
170
171
172def get_args() -> argparse.Namespace:
173    parser = argparse.ArgumentParser()
174
175    parser.add_argument("file", type=str, help=""" The file to format. """)
176    parser.add_argument(
177        "--start",
178        type=int,
179        default=None,
180        help=""" Optional line number to start formatting at. """,
181    )
182    parser.add_argument(
183        "--stop",
184        type=int,
185        default=None,
186        help=""" Optional line number to stop formatting at. """,
187    )
188    parser.add_argument(
189        "-nb",
190        "--noblack",
191        action="store_true",
192        help=""" Don't format file with Black after sorting. """,
193    )
194    parser.add_argument(
195        "-o",
196        "--output",
197        default=None,
198        help=""" Write changes to this file, otherwise changes are written back to the original file. """,
199    )
200    args = parser.parse_args()
201
202    return args
203
204
205def main(args: argparse.Namespace | None = None):
206    if not args:
207        args = get_args()
208    source = Pathier(args.file).read_text()
209    source = seat(source, args.start, args.stop)
210    if not args.noblack:
211        source = black.format_str(source, mode=black.Mode())
212    Pathier(args.output or args.file).write_text(source)
213
214
215if __name__ == "__main__":
216    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 Order:
29class Order:
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
39    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
40        return sorted(nodes, key=lambda node: node.name)
41
42    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
43        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
44        dunders = self.sort_nodes_by_name(dunders)
45        init = None
46        for i, dunder in enumerate(dunders):
47            if dunder.name == "__init__":
48                init = dunders.pop(i)
49                break
50        if init:
51            dunders.insert(0, init)
52        return dunders
53
54    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
55        """Sort assignment statments."""
56
57        def get_name(node: ast.stmt) -> str:
58            type_ = type(node)
59            if type_ == ast.Assign:
60                return node.targets[0].id
61            else:
62                return node.target.id
63
64        return sorted(assigns, key=get_name)
65
66    def sort(self) -> list[ast.stmt]:
67        """Sort and return members as a single list."""
68        self.dunders = self.sort_dunders(self.dunders)
69        self.functions = self.sort_nodes_by_name(self.functions)
70        self.properties = self.sort_nodes_by_name(self.properties)
71        self.assigns = self.sort_assigns(self.assigns)
72        return (
73            self.before
74            + self.assigns
75            + self.dunders
76            + self.properties
77            + self.functions
78            + self.seats
79            + self.after
80        )
Order()
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 = []
def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
39    def sort_nodes_by_name(self, nodes: list[ast.stmt]) -> list[ast.stmt]:
40        return sorted(nodes, key=lambda node: node.name)
def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
42    def sort_dunders(self, dunders: list[ast.stmt]) -> list[ast.stmt]:
43        """Sort `dunders` alphabetically, except `__init__` is placed at the front, if it exists."""
44        dunders = self.sort_nodes_by_name(dunders)
45        init = None
46        for i, dunder in enumerate(dunders):
47            if dunder.name == "__init__":
48                init = dunders.pop(i)
49                break
50        if init:
51            dunders.insert(0, init)
52        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]:
54    def sort_assigns(self, assigns: list[ast.stmt]) -> list[ast.stmt]:
55        """Sort assignment statments."""
56
57        def get_name(node: ast.stmt) -> str:
58            type_ = type(node)
59            if type_ == ast.Assign:
60                return node.targets[0].id
61            else:
62                return node.target.id
63
64        return sorted(assigns, key=get_name)

Sort assignment statments.

def sort(self) -> list[ast.stmt]:
66    def sort(self) -> list[ast.stmt]:
67        """Sort and return members as a single list."""
68        self.dunders = self.sort_dunders(self.dunders)
69        self.functions = self.sort_nodes_by_name(self.functions)
70        self.properties = self.sort_nodes_by_name(self.properties)
71        self.assigns = self.sort_assigns(self.assigns)
72        return (
73            self.before
74            + self.assigns
75            + self.dunders
76            + self.properties
77            + self.functions
78            + self.seats
79            + self.after
80        )

Sort and return members as a single list.

def seat( source: str, start_line: int | None = None, stop_line: int | None = None) -> str:
 83def seat(
 84    source: str, start_line: int | None = None, stop_line: int | None = None
 85) -> str:
 86    """Sort the contents of classes in `source`, where `source` is parsable Python code.
 87    Anything not inside a class will be untouched.
 88
 89    The modified `source` will be returned.
 90
 91    #### :params:
 92
 93    * `start_line`: Only sort contents after this line.
 94
 95    * `stop_line`: Only sort contents before this line.
 96
 97    If you have class contents that are grouped a certain way and you want the groups individually sorted
 98    so that the grouping is maintained, you can use `# Seat` to demarcate the groups.
 99
100    Note: Comments that are in a class body, but not in a function will remain at the same line.
101
102    i.e. if the source is:
103    >>> class MyClass():
104    >>>     {arbitrary lines of code}
105    >>>     # Seat
106    >>>     {more arbitrary code}
107    >>>     # Seat
108    >>>     {yet more code}
109
110    Then the three sets of code in brackets will be sorted independently from one another
111    (assuming no values are given for `start_line` or `stop_line`).
112
113    #### :Sorting and Priority:
114
115    * Class variables declared in class body outside of a function
116    * Dunder methods
117    * Functions decorated with `property` or corresponding `.setter` and `.deleter` methods
118    * Class functions
119
120    Each of these groups will be sorted alphabetically with respect to themselves.
121
122    The only exception is for dunder methods.
123    They will be sorted alphabetically except that `__init__` will be first.
124    """
125    tree = ast.parse(source, type_comments=True)
126    start_line = start_line or 0
127    stop_line = stop_line or len(source.splitlines()) + 1
128    sections = get_seat_sections(source)
129    comments = []
130    for section in sections:
131        for i, stmt in enumerate(tree.body):
132            if type(stmt) == ast.ClassDef:
133                order = Order()
134                for child in stmt.body:
135                    type_ = type(child)
136                    if child.lineno <= start_line or child.lineno < section[0]:
137                        order.before.append(child)
138                    elif stop_line < child.lineno or child.lineno > section[1]:
139                        order.after.append(child)
140                    elif type_ == ast.Comment:
141                        if "# Seat" in child.value:
142                            order.seats.append(child)
143                        else:
144                            comments.append(child)
145                    elif type_ in [ast.Assign, ast.AugAssign, ast.AnnAssign]:
146                        order.assigns.append(child)
147                    elif child.name.startswith("__") and child.name.endswith("__"):
148                        order.dunders.append(child)
149                    elif child.decorator_list:
150                        for decorator in child.decorator_list:
151                            decorator_type = type(decorator)
152                            if (
153                                decorator_type == ast.Name
154                                and decorator.id == "property"
155                            ) or (
156                                decorator_type == ast.Attribute
157                                and decorator.attr in ["setter", "deleter"]
158                            ):
159                                order.properties.append(child)
160                                break
161                        if child not in order.properties:
162                            order.functions.append(child)
163                    else:
164                        order.functions.append(child)
165                tree.body[i].body = order.sort()
166    # Put comments back in
167    source = ast.unparse(tree).splitlines()
168    for comment in comments:
169        source.insert(comment.lineno - 1, f"{' '*comment.col_offset}{comment.value}")
170    return "\n".join(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.

Note: Comments that are in a class body, but not in a function will remain at the same line.

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:
173def get_args() -> argparse.Namespace:
174    parser = argparse.ArgumentParser()
175
176    parser.add_argument("file", type=str, help=""" The file to format. """)
177    parser.add_argument(
178        "--start",
179        type=int,
180        default=None,
181        help=""" Optional line number to start formatting at. """,
182    )
183    parser.add_argument(
184        "--stop",
185        type=int,
186        default=None,
187        help=""" Optional line number to stop formatting at. """,
188    )
189    parser.add_argument(
190        "-nb",
191        "--noblack",
192        action="store_true",
193        help=""" Don't format file with Black after sorting. """,
194    )
195    parser.add_argument(
196        "-o",
197        "--output",
198        default=None,
199        help=""" Write changes to this file, otherwise changes are written back to the original file. """,
200    )
201    args = parser.parse_args()
202
203    return args
def main(args: argparse.Namespace | None = None):
206def main(args: argparse.Namespace | None = None):
207    if not args:
208        args = get_args()
209    source = Pathier(args.file).read_text()
210    source = seat(source, args.start, args.stop)
211    if not args.noblack:
212        source = black.format_str(source, mode=black.Mode())
213    Pathier(args.output or args.file).write_text(source)