diff --git a/utils.py b/utils.py index 19748af..63d7c2a 100644 --- a/utils.py +++ b/utils.py @@ -1,16 +1,15 @@ import re -_CLASS_FINDER = re.compile(r'^\s*class\s+([A-Z][a-zA-Z0-9_]*)\s*\((?:ABC|object)\)\s*:\s*') +_CLASS_FINDER = re.compile(r'^\s*class\s+([A-Z][a-zA-Z0-9_]*)(?:\s*\([^)]*\))?\s*:') def guess_classname(code: str) -> str: for line in code.splitlines(): if match := _CLASS_FINDER.search(line): return match.group(1) - raise ValueError('Cannot find classname in code. Expected a class definition like `class ClassName(ABC):`') + raise ValueError('Cannot find classname in code. Expected a class definition like `class ClassName:`') def camel_to_snake(input: str) -> str: s1 = re.sub(r'(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])', '_', input) - s2 = re.sub(r'([a-zA-Z])([0-9])', r'\1\2', s1) - return re.sub(r'([0-9])([A-Z])', r'\1_\2', s2).lower() + return re.sub(r'([0-9])([A-Z])', r'\1_\2', s1).lower() diff --git a/utils_test.py b/utils_test.py index a040080..9d88536 100644 --- a/utils_test.py +++ b/utils_test.py @@ -27,8 +27,7 @@ def test_guess_classname_with_multiple_classes(self): def test_guess_classname_with_complex_class_definition(self): code = "class MyComplexClass(object, metaclass=MyMeta): pass" - with self.assertRaises(ValueError): - guess_classname(code) + self.assertEqual(guess_classname(code), "MyComplexClass") def test_camel_to_snake_edge_case(self): self.assertEqual(camel_to_snake(""), "") @@ -52,3 +51,33 @@ def test_camel_to_snake_with_numbers(self): def test_camel_to_snake_single_word(self): self.assertEqual(camel_to_snake("hello"), "hello") self.assertEqual(camel_to_snake("WORLD"), "world") + + def test_guess_classname_no_parentheses(self): + code = "class MyClass: pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_guess_classname_other_parent(self): + code = "class MyClass(Base): pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_guess_classname_multiple_parents(self): + code = "class MyClass(Base1, Base2): pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_guess_classname_with_comment(self): + code = "class MyClass(ABC): # This is a class\n pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_guess_classname_after_comments(self): + code = "# Comment\n\nclass MyClass(ABC): pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_camel_to_snake_already_snake(self): + self.assertEqual(camel_to_snake("already_snake_case"), "already_snake_case") + + def test_camel_to_snake_mixed_case(self): + self.assertEqual(camel_to_snake("My_Mixed_Case"), "my_mixed_case") + + def test_camel_to_snake_complex_numbers(self): + self.assertEqual(camel_to_snake("A1B2C3"), "a1_b2_c3") + self.assertEqual(camel_to_snake("v6Address"), "v6_address")