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