From a07ae32a72d074b91767e806ae53e1ffe8d7c933 Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Sun, 28 Jul 2024 14:18:17 -0400
Subject: [PATCH] use a Protocol for ast_to_offset

---
 add_trailing_comma/_ast_helpers.py     | 10 +++++++++-
 add_trailing_comma/_plugins/calls.py   |  2 +-
 add_trailing_comma/_plugins/classes.py |  2 +-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/add_trailing_comma/_ast_helpers.py b/add_trailing_comma/_ast_helpers.py
index 09759ee..942803e 100644
--- a/add_trailing_comma/_ast_helpers.py
+++ b/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)
diff --git a/add_trailing_comma/_plugins/calls.py b/add_trailing_comma/_plugins/calls.py
index bd707dd..db6380c 100644
--- a/add_trailing_comma/_plugins/calls.py
+++ b/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)
diff --git a/add_trailing_comma/_plugins/classes.py b/add_trailing_comma/_plugins/classes.py
index 0d423e9..2e4af74 100644
--- a/add_trailing_comma/_plugins/classes.py
+++ b/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: