reverted decorator update to autopep8 lib

now the default test decorator will be @mark.testomatio('id') and from pytest import mark will be added to every test file

README.md

@@ -100,6 +100,8 @@ You would need to decide when you want to upload your test artifacts to cloud st
 ```python
 # content of conftest.py
 import pytest
+import random
+import os
 from typing import Dict
 from pytest import StashKey, CollectReport
 from playwright.sync_api import Page
@@ -156,10 +158,10 @@ To make the experience more consistent, it uses standard pytest markers.
 testomat.io test id is a string value that starts with `@T` and has 8 symbols after.
 
 ```python
-import pytest
+from pytest import mark
 
 
-@pytest.mark.testomatio('@T96c700e6')
+@mark.testomatio('@T96c700e6')
 def test_example():
     assert 2 + 2 == 4
 ```

pyproject.toml

@@ -13,7 +13,7 @@ version_provider = "pep621"
 update_changelog_on_bump = true
 [project]
 name = "pytestomatio"
-version = "2.2.3"
+version = "2.2.4"
 
 dependencies = [
     "requests>=2.29.0",

pytestomatio/decorator_updater.py

@@ -1,8 +1,11 @@
-import libcst as cst
-from typing import List, Tuple, Union
+import ast
+import autopep8
 
-class DecoratorUpdater(cst.CSTTransformer):
-    def __init__(self, mapped_tests: List[Tuple[str, int]], all_tests: List[str], decorator_name: str):
+pytest_mark = 'pytest', 'mark'
+
+
+class DecoratorUpdater(ast.NodeTransformer):
+    def __init__(self, mapped_tests: list[tuple[str, int]], all_tests: list[str], decorator_name: str):
         self.mapped_tests = mapped_tests
         self.all_tests = all_tests
         self.decorator_name = decorator_name
@@ -12,70 +15,65 @@ def _get_id_by_title(self, title: str):
             if pair[0] == title:
                 return pair[1]
 
-    def _remove_decorator(self, node: cst.FunctionDef) -> cst.FunctionDef:
+    def _remove_decorator(self, node: ast.FunctionDef) -> ast.FunctionDef:
         node.decorator_list = [decorator for decorator in node.decorator_list if
-                               not (isinstance(decorator, cst.Call) and decorator.func.attr == self.decorator_name)]
+                               not (isinstance(decorator, ast.Call) and decorator.func.attr == self.decorator_name)]
         return node
 
-    def remove_decorators(self, tree: cst.Module) -> cst.Module:
-        for node in cst.walk(tree):
-            if isinstance(node, cst.FunctionDef):
+    def remove_decorators(self, tree: ast.Module) -> ast.Module:
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef):
                 self.visit_FunctionDef(node, remove=True)
         return tree
 
-    def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef:
-        if original_node.name.value in self.all_tests:
-            test_id = self._get_id_by_title(original_node.name.value)
-            if test_id is None:
-                return updated_node
-
-            deco_name = f'pytest.mark.{self.decorator_name}("{test_id}")'
-            decorator = cst.Decorator(decorator=cst.parse_expression(deco_name))
-
-            # Check if the decorator already exists
-            for existing_decorator in original_node.decorators:
-                if isinstance(existing_decorator.decorator, cst.Call) and \
-                isinstance(existing_decorator.decorator.func, cst.Attribute) and \
-                existing_decorator.decorator.func.attr.value == self.decorator_name:
-                    # The decorator already exists, so we don't add it
-                    return updated_node
-
-            # The decorator doesn't exist, so we add it
-            return updated_node.with_changes(decorators=[decorator] + list(updated_node.decorators))
-        return updated_node
+    def visit_FunctionDef(self, node: ast.FunctionDef, remove=False) -> ast.FunctionDef:
+        if remove:
+            return self._remove_decorator(node)
+        else:
+            if node.name in self.all_tests:
+                if not any(isinstance(decorator, ast.Call) and
+                           decorator.func.attr == self.decorator_name
+                           for decorator in node.decorator_list):
+                    test_id = self._get_id_by_title(node.name)
+                    deco_name = f'mark.{self.decorator_name}(\'{test_id}\')'
+                    decorator = ast.Name(id=deco_name, ctx=ast.Load())
+                    node.decorator_list = [decorator] + node.decorator_list
+        return node
 
-class DecoratorRemover(cst.CSTTransformer):
-    def __init__(self, decorator_name: str):
-        self.decorator_name = decorator_name
+    def insert_pytest_mark_import(self, tree: ast.Module, module_name: str, decorator_name: str) -> None:
+        # Check if the import statement already exists
+        if not any(
+                isinstance(node, ast.ImportFrom) and
+                node.module == module_name and
+                any(alias.name == decorator_name for alias in node.names)
+                for node in tree.body
+        ):
+            import_node = ast.ImportFrom(
+                module=module_name,
+                names=[ast.alias(name=decorator_name, asname=None)],
+                level=0
+            )
+            tree.body.insert(0, import_node)
 
-    def leave_Decorator(self, original_node: cst.Decorator, updated_node: cst.Decorator) -> Union[cst.Decorator, cst.RemovalSentinel]:
-        if isinstance(original_node.decorator, cst.Call) and \
-           isinstance(original_node.decorator.func, cst.Attribute) and \
-           original_node.decorator.func.attr.value == self.decorator_name and \
-           isinstance(original_node.decorator.func.value, cst.Attribute) and \
-           original_node.decorator.func.value.attr.value == 'mark' and \
-           isinstance(original_node.decorator.func.value.value, cst.Name) and \
-           original_node.decorator.func.value.value.value == 'pytest':
-            return cst.RemovalSentinel.REMOVE
-        return updated_node
 
 def update_tests(file: str,
-                 mapped_tests: List[Tuple[str, int]],
-                 all_tests: List[str],
+                 mapped_tests: list[tuple[str, int]],
+                 all_tests: list[str],
                  decorator_name: str,
                  remove=False):
     with open(file, 'r') as f:
         source_code = f.read()
 
-    tree = cst.parse_module(source_code)
+    tree = ast.parse(source_code)
     transform = DecoratorUpdater(mapped_tests, all_tests, decorator_name)
     if remove:
-        transform = DecoratorRemover(decorator_name)
-        tree = tree.visit(transform)
+        transform.remove_decorators(tree)
     else:
-        transform = DecoratorUpdater(mapped_tests, all_tests, decorator_name)
-        tree = tree.visit(transform)
-    updated_source_code = tree.code
+        tree = transform.visit(tree)
+        transform.insert_pytest_mark_import(tree, *pytest_mark)
+    updated_source_code = ast.unparse(tree)
+
+    pep8_source_code = autopep8.fix_code(updated_source_code)
 
     with open(file, "w") as file:
-        file.write(updated_source_code)
+        file.write(pep8_source_code)

tests/test_decorators.py

@@ -0,0 +1,6 @@
+from pytest import mark
+
+
+@mark.testomatio('@T741f0586')
+def test_something():
+    assert 1 == 1