diff --git a/flake8_simplify.py b/flake8_simplify.py index ecb68b0..f730848 100644 --- a/flake8_simplify.py +++ b/flake8_simplify.py @@ -1093,24 +1093,41 @@ def _get_sim119(node: ast.ClassDef) -> List[Tuple[int, int, str]]: "__repr__", "__str__", ] - has_only_constructur_function = True + has_only_dataclass_functions = True + has_any_functions = False for body_el in node.body: - if ( - isinstance(body_el, ast.FunctionDef) - and body_el.name not in dataclass_functions - ): - has_only_constructur_function = False - break + if isinstance(body_el, (ast.FunctionDef, ast.AsyncFunctionDef)): + has_any_functions = True + if body_el.name == "__init__": + # Ensure constructor only has pure assignments + # without any calculation. + has_complex_statements = False + for el in body_el.body: + if not isinstance(el, ast.Assign): + has_complex_statements = True + break + else: + # It is an assignment, but we only allow + # `self.attribute = name`. + if any( + [ + not isinstance(target, ast.Attribute) + for target in el.targets + ] + ) or not isinstance(el.value, ast.Name): + has_complex_statements = True + break + if body_el.name not in dataclass_functions: + has_only_dataclass_functions = False - if not ( - has_only_constructur_function - and sum(1 for el in node.body if isinstance(el, ast.FunctionDef)) > 0 + if ( + has_any_functions + and has_only_dataclass_functions + and not has_complex_statements ): - return errors - - errors.append( - (node.lineno, node.col_offset, SIM119.format(classname=node.name)) - ) + errors.append( + (node.lineno, node.col_offset, SIM119.format(classname=node.name)) + ) return errors diff --git a/tests/test_simplify.py b/tests/test_simplify.py index 1f791a7..24ba5dd 100644 --- a/tests/test_simplify.py +++ b/tests/test_simplify.py @@ -476,6 +476,25 @@ def test_sim118_del_key(): def test_sim119(): results = _results( """ +class FooBar: + def __init__(self, a, b): + self.a = a + self.b = b +""" + ) + assert results == {"2:0 SIM119 Use a dataclass for 'class FooBar'"} + + +def test_sim119_ignored_dunder_methods(): + """ + Dunder methods do not make a class not be a dataclass candidate. + Examples for dunder (double underscore) methods are: + * __str__ + * __eq__ + * __hash__ + """ + results = _results( + """ class FooBar: def __init__(self, a, b): self.a = a @@ -488,6 +507,32 @@ def __str__(self): assert results == {"2:0 SIM119 Use a dataclass for 'class FooBar'"} +def test_sim119_async(): + results = _results( + """ +class FooBar: + def __init__(self, a, b): + self.a = a + self.b = b + + async def foo(self): + return "FooBar" +""" + ) + assert results == set() + + +def test_sim119_constructor_processing(): + results = _results( + """ +class FooBar: + def __init__(self, a): + self.a = a + 5 +""" + ) + assert results == set() + + def test_sim119_pydantic(): results = _results( """