.

Author: vishalshenoy

run_format_example.py

@@ -0,0 +1,89 @@
+import codegen
+from codegen import Codebase
+from codegen.sdk.core.detached_symbols.function_call import FunctionCall
+
+
+@codegen.function("useSuspenseQuery-to-useSuspenseQueries")
+def run(codebase: Codebase):
+    """Convert useSuspenseQuery calls to useSuspenseQueries in a React codebase.
+
+    This codemod:
+    1. Finds all files containing useSuspenseQuery
+    2. Adds the necessary import statement
+    3. Converts multiple useSuspenseQuery calls to a single useSuspenseQueries call
+    """
+    # Import statement for useSuspenseQueries
+    import_str = "import { useQuery, useSuspenseQueries } from '@tanstack/react-query'"
+
+    # Track statistics
+    files_modified = 0
+    functions_modified = 0
+
+    # Iterate through all files in the codebase
+    for file in codebase.files:
+        if "useSuspenseQuery" not in file.source:
+            continue
+
+        print(f"Processing {file.filepath}")
+        # Add the import statement
+        file.add_import_from_import_string(import_str)
+        file_modified = False
+
+        # Iterate through all functions in the file
+        for function in file.functions:
+            if "useSuspenseQuery" not in function.source:
+                continue
+
+            results = []  # Store left-hand side of assignments
+            queries = []  # Store query arguments
+            old_statements = []  # Track statements to replace
+
+            # Find useSuspenseQuery assignments
+            for stmt in function.code_block.assignment_statements:
+                if not isinstance(stmt.right, FunctionCall):
+                    continue
+
+                fcall = stmt.right
+                if fcall.name != "useSuspenseQuery":
+                    continue
+
+                old_statements.append(stmt)
+                results.append(stmt.left.source)
+                queries.append(fcall.args[0].value.source)
+
+            # Convert to useSuspenseQueries if needed
+            if old_statements:
+                new_query = f"const [{', '.join(results)}] = useSuspenseQueries({{queries: [{', '.join(queries)}]}})"
+                print(
+                    f"Converting useSuspenseQuery to useSuspenseQueries in {function.name}"
+                )
+
+                # Print the diff
+                print("\nOriginal code:")
+                for stmt in old_statements:
+                    print(f"- {stmt.source}")
+                print("\nNew code:")
+                print(f"+ {new_query}")
+                print("-" * 50)
+
+                # Replace old statements with new query
+                for stmt in old_statements:
+                    stmt.edit(new_query)
+
+                functions_modified += 1
+                file_modified = True
+
+        if file_modified:
+            files_modified += 1
+
+    print("\nModification complete:")
+    print(f"Files modified: {files_modified}")
+    print(f"Functions modified: {functions_modified}")
+
+
+if __name__ == "__main__":
+    print("Initializing codebase...")
+    codebase = Codebase.from_repo("deepfence/ThreatMapper")
+
+    print("Running codemod...")
+    run(codebase)

sqlalchemy_1.6_to_2.0/input_repo/database.py

@@ -3,7 +3,9 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
 
-SQLALCHEMY_DATABASE_URL = "postgresql://user:password@localhost/dbname"  # Change to your database URL
+SQLALCHEMY_DATABASE_URL = (
+    "postgresql://user:password@localhost/dbname"  # Change to your database URL
+)
 
 engine = create_engine(SQLALCHEMY_DATABASE_URL)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

sqlalchemy_1.6_to_2.0/input_repo/models.py

@@ -2,6 +2,7 @@
 from sqlalchemy.orm import relationship, backref
 from database import Base
 
+
 class Publisher(Base):
     __tablename__ = "publishers"