use a Protocol for ast_to_offset

Author: asottile

add_trailing_comma/_ast_helpers.py

@@ -2,6 +2,7 @@
 
 import ast
 import warnings
+from typing import Protocol
 
 from tokenize_rt import Offset
 
@@ -12,5 +13,12 @@ def ast_parse(contents_text: str) -> ast.Module:
         return ast.parse(contents_text.encode())
 
 
-def ast_to_offset(node: ast.AST) -> Offset:
+class _HasOffsetInfo(Protocol):
+    @property
+    def lineno(self) -> int: ...
+    @property
+    def col_offset(self) -> int: ...
+
+
+def ast_to_offset(node: _HasOffsetInfo) -> Offset:
     return Offset(node.lineno, node.col_offset)

add_trailing_comma/_plugins/calls.py

@@ -34,7 +34,7 @@ def visit_Call(
         state: State,
         node: ast.Call,
 ) -> Iterable[tuple[Offset, TokenFunc]]:
-    argnodes = [*node.args, *node.keywords]
+    argnodes: list[ast.expr | ast.keyword] = [*node.args, *node.keywords]
     arg_offsets = set()
     for argnode in argnodes:
         offset = ast_to_offset(argnode)

add_trailing_comma/_plugins/classes.py

@@ -37,7 +37,7 @@ def visit_ClassDef(
     # starargs are allowed in py3 class definitions, py35+ allows trailing
     # commas.  py34 does not, but adding an option for this very obscure
     # case seems not worth it.
-    args = [*node.bases, *node.keywords]
+    args: list[ast.expr | ast.keyword] = [*node.bases, *node.keywords]
     arg_offsets = {ast_to_offset(arg) for arg in args}
 
     if arg_offsets: