Remove unused functions, adjust the tests for the new ones

Author: arnaud-ma

src/tap/utils.py

@@ -185,17 +185,12 @@ def is_positional_arg(*name_or_flags) -> bool:
     return not is_option_arg(*name_or_flags)
 
 
-def _tokenize_source(source: str) -> Generator[tokenize.TokenInfo, None, None]:
+def tokenize_source(source: str) -> Generator[tokenize.TokenInfo, None, None]:
     """Returns a generator for the tokens of the object's source code, given the source code."""
     return tokenize.generate_tokens(StringIO(source).readline)
 
 
-def tokenize_source(obj: object) -> Generator:
-    """Returns a generator for the tokens of the object's source code."""
-    return _tokenize_source(inspect.getsource(obj))
-
-
-def _get_class_column(tokens: Iterable[tokenize.TokenInfo]) -> int:
+def get_class_column(tokens: Iterable[tokenize.TokenInfo]) -> int:
     """Determines the column number for class variables in a class, given the tokens of the class."""
     first_line = 1
     for token_type, token, (start_line, start_column), (end_line, end_column), line in tokens:
@@ -205,14 +200,10 @@ def _get_class_column(tokens: Iterable[tokenize.TokenInfo]) -> int:
             continue
 
         return start_column
+    raise ValueError("Could not find any class variables in the class.")
 
 
-def get_class_column(cls: type) -> int:
-    """Determines the column number for class variables in a class."""
-    return _get_class_column(tokenize_source(cls))
-
-
-def _source_line_to_tokens(tokens: Iterable[tokenize.TokenInfo]) -> Dict[int, List[Dict[str, Union[str, int]]]]:
+def source_line_to_tokens(tokens: Iterable[tokenize.TokenInfo]) -> Dict[int, List[Dict[str, Union[str, int]]]]:
     """
     Gets a dictionary mapping from line number to a dictionary of tokens on that line for an object's source code,
     given the tokens of the object's source code.
@@ -232,12 +223,7 @@ def _source_line_to_tokens(tokens: Iterable[tokenize.TokenInfo]) -> Dict[int, Li
     return line_to_tokens
 
 
-def source_line_to_tokens(obj: object) -> Dict[int, List[Dict[str, Union[str, int]]]]:
-    """Gets a dictionary mapping from line number to a dictionary of tokens on that line for an object's source code."""
-    return _source_line_to_tokens(tokenize_source(obj))
-
-
-def _get_subsequent_assign_lines(source_cls: str) -> Set[int]:
+def get_subsequent_assign_lines(source_cls: str) -> Set[int]:
     """
     For all multiline assign statements, get the line numbers after the first line of the assignment,
     given the source code of the object.
@@ -283,25 +269,22 @@ def _get_subsequent_assign_lines(source_cls: str) -> Set[int]:
 
     return assign_lines
 
-def get_subsequent_assign_lines(cls: type) -> Set[int]:
-    """For all multiline assign statements, get the line numbers after the first line of the assignment."""
-    return _get_subsequent_assign_lines(inspect.getsource(cls))
 
 def get_class_variables(cls: type) -> Dict[str, Dict[str, str]]:
     """Returns a dictionary mapping class variables to their additional information (currently just comments)."""
     # Get the source code and tokens of the class
     source_cls = inspect.getsource(cls)
-    tokens = tuple(_tokenize_source(source_cls))
+    tokens = tuple(tokenize_source(source_cls))
 
     # Get mapping from line number to tokens
-    line_to_tokens = _source_line_to_tokens(tokens)
+    line_to_tokens = source_line_to_tokens(tokens)
 
     # Get class variable column number
-    class_variable_column = _get_class_column(tokens)
+    class_variable_column = get_class_column(tokens)
 
     # For all multiline assign statements, get the line numbers after the first line of the assignment
     # This is used to avoid identifying comments in multiline assign statements
-    subsequent_assign_lines = _get_subsequent_assign_lines(source_cls)
+    subsequent_assign_lines = get_subsequent_assign_lines(source_cls)
 
     # Extract class variables
     class_variable = None

tests/test_utils.py

@@ -1,4 +1,5 @@
 from argparse import ArgumentTypeError
+import inspect
 import json
 import os
 import subprocess
@@ -11,6 +12,7 @@
     get_class_column,
     get_class_variables,
     GitInfo,
+    tokenize_source,
     type_to_str,
     get_literals,
     TupleTypeEnforcer,
@@ -145,7 +147,8 @@ def test_column_simple(self):
         class SimpleColumn:
             arg = 2
 
-        self.assertEqual(get_class_column(SimpleColumn), 12)
+        tokens = tokenize_source(inspect.getsource(SimpleColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
     def test_column_comment(self):
         class CommentColumn:
@@ -158,28 +161,32 @@ class CommentColumn:
 
             arg = 2
 
-        self.assertEqual(get_class_column(CommentColumn), 12)
+        tokens = tokenize_source(inspect.getsource(CommentColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
     def test_column_space(self):
         class SpaceColumn:
 
             arg = 2
 
-        self.assertEqual(get_class_column(SpaceColumn), 12)
+        tokens = tokenize_source(inspect.getsource(SpaceColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
     def test_column_method(self):
         class FuncColumn:
             def func(self):
                 pass
 
-        self.assertEqual(get_class_column(FuncColumn), 12)
+        tokens = tokenize_source(inspect.getsource(FuncColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
     def test_dataclass(self):
         @class_decorator
         class DataclassColumn:
             arg: int = 5
 
-        self.assertEqual(get_class_column(DataclassColumn), 12)
+        tokens = tokenize_source(inspect.getsource(DataclassColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
     def test_dataclass_method(self):
         def wrapper(f):
@@ -191,7 +198,8 @@ class DataclassColumn:
             def func(self):
                 pass
 
-        self.assertEqual(get_class_column(DataclassColumn), 12)
+        tokens = tokenize_source(inspect.getsource(DataclassColumn))
+        self.assertEqual(get_class_column(tokens), 12)
 
 
 class ClassVariableTests(TestCase):