diff --git a/integration_tests/test_django_model_without_dunder_str.py b/integration_tests/test_django_model_without_dunder_str.py new file mode 100644 index 00000000..b059b1a2 --- /dev/null +++ b/integration_tests/test_django_model_without_dunder_str.py @@ -0,0 +1,52 @@ +from core_codemods.django_model_without_dunder_str import ( + DjangoModelWithoutDunderStr, + DjangoModelWithoutDunderStrTransformer, +) +from integration_tests.base_test import ( + BaseIntegrationTest, + original_and_expected_from_code_path, +) + + +class TestDjangoModelWithoutDunderStr(BaseIntegrationTest): + codemod = DjangoModelWithoutDunderStr + code_path = "tests/samples/django-project/mysite/mysite/models.py" + original_code, expected_new_code = original_and_expected_from_code_path( + code_path, + [ + (15, """\n"""), + (16, """ def __str__(self):\n"""), + (17, """ model_name = self.__class__.__name__\n"""), + ( + 18, + """ fields_str = ", ".join([f"{field.name}={getattr(self, field.name)}" for field in self._meta.fields])\n""", + ), + (19, """ return f"{model_name}({fields_str})"\n"""), + ], + ) + + # fmt: off + expected_diff =( + """--- \n""" + """+++ \n""" + """@@ -11,3 +11,8 @@\n""" + """ content = models.CharField(max_length=200)\n""" + """ class Meta:\n""" + """ app_label = 'myapp'\n""" + """+\n""" + """+ def __str__(self):\n""" + """+ model_name = self.__class__.__name__\n""" + """+ fields_str = ", ".join([f"{field.name}={getattr(self, field.name)}" for field in self._meta.fields])\n""" + """+ return f"{model_name}({fields_str})"\n""" + ) + # fmt: on + + expected_line_change = "9" + change_description = DjangoModelWithoutDunderStrTransformer.change_description + num_changed_files = 1 + + def check_code_after(self): + """Executes models.py and instantiates the model to ensure expected str representation""" + module = super().check_code_after() + inst = module.Message(pk=1, author="name", content="content") + assert str(inst) == "Message(id=1, author=name, content=content)" diff --git a/src/codemodder/codemods/test/integration_utils.py b/src/codemodder/codemods/test/integration_utils.py index 9e970f96..6caa5e73 100644 --- a/src/codemodder/codemods/test/integration_utils.py +++ b/src/codemodder/codemods/test/integration_utils.py @@ -8,6 +8,7 @@ from codemodder import __version__ from codemodder import registry from .validations import execute_code +from types import ModuleType SAMPLES_DIR = "tests/samples" # Enable import of test modules from test directory @@ -139,11 +140,13 @@ def check_code_before(self): code = f.read() assert code == self.original_code - def check_code_after(self): + def check_code_after(self) -> ModuleType: with open(self.code_path, "r", encoding="utf-8") as f: new_code = f.read() assert new_code == self.expected_new_code - execute_code(path=self.code_path, allowed_exceptions=self.allowed_exceptions) + return execute_code( + path=self.code_path, allowed_exceptions=self.allowed_exceptions + ) def test_file_rewritten(self): """ diff --git a/src/codemodder/codemods/test/validations.py b/src/codemodder/codemods/test/validations.py index 915569f5..eab41c7c 100644 --- a/src/codemodder/codemods/test/validations.py +++ b/src/codemodder/codemods/test/validations.py @@ -1,5 +1,7 @@ import importlib.util import tempfile +from types import ModuleType +from typing import Optional def execute_code(*, path=None, code=None, allowed_exceptions=None): @@ -11,20 +13,29 @@ def execute_code(*, path=None, code=None, allowed_exceptions=None): ), "Must pass either path to code or code as a str." if path: - _run_code(path, allowed_exceptions) - return + return _run_code(path, allowed_exceptions) with tempfile.NamedTemporaryFile(suffix=".py", mode="w+t") as temp: temp.write(code) - _run_code(temp.name, allowed_exceptions) + return _run_code(temp.name, allowed_exceptions) -def _run_code(path, allowed_exceptions=None): - """Execute the code in `path` in its own namespace.""" +def _run_code(path, allowed_exceptions=None) -> Optional[ModuleType]: + """ + Execute the code in `path` in its own namespace. + Return loaded module for any additional testing later on. + """ allowed_exceptions = allowed_exceptions or () spec = importlib.util.spec_from_file_location("output_code", path) + if not spec: + return None + module = importlib.util.module_from_spec(spec) + if not spec.loader: + return None try: spec.loader.exec_module(module) except allowed_exceptions: pass + + return module diff --git a/src/core_codemods/__init__.py b/src/core_codemods/__init__.py index 0129aecf..9f77b5d7 100644 --- a/src/core_codemods/__init__.py +++ b/src/core_codemods/__init__.py @@ -62,6 +62,7 @@ from .lazy_logging import LazyLogging from .str_concat_in_seq_literal import StrConcatInSeqLiteral from .fix_async_task_instantiation import FixAsyncTaskInstantiation +from .django_model_without_dunder_str import DjangoModelWithoutDunderStr registry = CodemodCollection( origin="pixee", @@ -118,6 +119,7 @@ LazyLogging, StrConcatInSeqLiteral, FixAsyncTaskInstantiation, + DjangoModelWithoutDunderStr, ], ) @@ -134,3 +136,4 @@ SonarDjangoJsonResponseType, ], ) +# Import and add Codemod class to registry above. diff --git a/src/core_codemods/django_model_without_dunder_str.py b/src/core_codemods/django_model_without_dunder_str.py index c65d6c5b..983b4223 100644 --- a/src/core_codemods/django_model_without_dunder_str.py +++ b/src/core_codemods/django_model_without_dunder_str.py @@ -37,16 +37,9 @@ def leave_ClassDef( self.report_change(original_node) - dunder_str = cst.FunctionDef( - leading_lines=[cst.EmptyLine(indent=False)], - name=cst.Name("__str__"), - params=cst.Parameters(params=[cst.Param(name=cst.Name("self"))]), - body=cst.IndentedBlock(body=[cst.SimpleStatementLine(body=[cst.Pass()])]), - ) new_body = updated_node.body.with_changes( - body=[*updated_node.body.body, dunder_str] + body=[*updated_node.body.body, dunder_str_method()] ) - return updated_node.with_changes(body=new_body) def implements_dunder_str(self, original_node: cst.ClassDef) -> bool: @@ -57,6 +50,24 @@ def implements_dunder_str(self, original_node: cst.ClassDef) -> bool: return False +def dunder_str_method() -> cst.FunctionDef: + self_body = cst.IndentedBlock( + body=[ + cst.parse_statement("model_name = self.__class__.__name__"), + cst.parse_statement( + 'fields_str = ", ".join([f"{field.name}={getattr(self, field.name)}" for field in self._meta.fields])' + ), + cst.parse_statement('return f"{model_name}({fields_str})"'), + ] + ) + return cst.FunctionDef( + leading_lines=[cst.EmptyLine(indent=False)], + name=cst.Name("__str__"), + params=cst.Parameters(params=[cst.Param(name=cst.Name("self"))]), + body=self_body, + ) + + DjangoModelWithoutDunderStr = CoreCodemod( metadata=Metadata( name="django-model-without-dunder-str", diff --git a/tests/codemods/test_django_model_without_dunder_str.py b/tests/codemods/test_django_model_without_dunder_str.py index 4c845002..73839318 100644 --- a/tests/codemods/test_django_model_without_dunder_str.py +++ b/tests/codemods/test_django_model_without_dunder_str.py @@ -48,7 +48,9 @@ def decorated_name(self): return f"***{self.name}***" def __str__(self): - pass + model_name = self.__class__.__name__ + fields_str = ", ".join([f"{field.name}={getattr(self, field.name)}" for field in self._meta.fields]) + return f"{model_name}({fields_str})" def something(): pass diff --git a/tests/samples/django-project/mysite/mysite/models.py b/tests/samples/django-project/mysite/mysite/models.py new file mode 100644 index 00000000..2086168c --- /dev/null +++ b/tests/samples/django-project/mysite/mysite/models.py @@ -0,0 +1,13 @@ +import django +from django.conf import settings +from django.db import models +# required to run this module standalone for testing +settings.configure() +django.setup() + + +class Message(models.Model): + author = models.CharField(max_length=100) + content = models.CharField(max_length=200) + class Meta: + app_label = 'myapp'