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())
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)]
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 )
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.
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.
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.
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.
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
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)