diff --git a/utils.py b/utils.py index 63d7c2a..cdafe2d 100644 --- a/utils.py +++ b/utils.py @@ -1,12 +1,11 @@ import re -_CLASS_FINDER = re.compile(r'^\s*class\s+([A-Z][a-zA-Z0-9_]*)(?:\s*\([^)]*\))?\s*:') +_CLASS_FINDER = re.compile(r'^\s*class\s+([A-Z][a-zA-Z0-9_]*)(?:\s|\(|:)', re.MULTILINE) def guess_classname(code: str) -> str: - for line in code.splitlines(): - if match := _CLASS_FINDER.search(line): - return match.group(1) + if match := _CLASS_FINDER.search(code): + return match.group(1) raise ValueError('Cannot find classname in code. Expected a class definition like `class ClassName:`') diff --git a/utils_test.py b/utils_test.py index 9d88536..ee9cbd5 100644 --- a/utils_test.py +++ b/utils_test.py @@ -81,3 +81,53 @@ def test_camel_to_snake_mixed_case(self): 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") + + def test_camel_to_snake_trailing_capital(self): + self.assertEqual(camel_to_snake("MyClassA"), "my_class_a") + + def test_camel_to_snake_all_caps(self): + self.assertEqual(camel_to_snake("HTTP"), "http") + self.assertEqual(camel_to_snake("HTTPServer"), "http_server") + + def test_guess_classname_indented(self): + code = " class IndentedClass:\n pass" + self.assertEqual(guess_classname(code), "IndentedClass") + + def test_guess_classname_with_multiline_docstring_before(self): + code = '"""\nThis is a docstring\n"""\nclass MyClass:\n pass' + self.assertEqual(guess_classname(code), "MyClass") + + def test_guess_classname_multiline_definition(self): + # Test that it works with multiline class definitions + code = "class MyClass(\n Base\n):\n pass" + self.assertEqual(guess_classname(code), "MyClass") + + def test_camel_to_snake_with_middle_acronym(self): + self.assertEqual(camel_to_snake("APIRequestResponse"), "api_request_response") + + def test_camel_to_snake_with_numbers_and_acronyms(self): + self.assertEqual(camel_to_snake("HTML5Parser"), "html5_parser") + self.assertEqual(camel_to_snake("JSON2HTML"), "json2_html") + + def test_camel_to_snake_starting_with_number(self): + self.assertEqual(camel_to_snake("123ClassName"), "123_class_name") + + def test_camel_to_snake_already_snake_with_numbers(self): + self.assertEqual(camel_to_snake("already_snake_123"), "already_snake_123") + + def test_camel_to_snake_with_special_characters(self): + # camel_to_snake doesn't explicitly handle dashes, but let's see what it does + self.assertEqual(camel_to_snake("StringWith-Dash"), "string_with-dash") + + def test_guess_classname_with_underscore(self): + code = "class My_Class: pass" + self.assertEqual(guess_classname(code), "My_Class") + + def test_guess_classname_not_at_start_of_line(self): + code = "# class NotMe:\nclass Me: pass" + self.assertEqual(guess_classname(code), "Me") + + def test_guess_classname_commented_out(self): + code = "# class NotMe: pass" + with self.assertRaises(ValueError): + guess_classname(code)