From 46c0d75b33d673e2d9da6560eb712514bc7b7bb6 Mon Sep 17 00:00:00 2001 From: Vishal Shenoy Date: Tue, 28 Jan 2025 15:03:20 -0800 Subject: [PATCH 1/2] . --- examples/dict_to_schema/README.md | 93 +++++++++++++++++++++++++++ examples/dict_to_schema/run.py | 102 ++++++++++++++++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 examples/dict_to_schema/README.md create mode 100644 examples/dict_to_schema/run.py diff --git a/examples/dict_to_schema/README.md b/examples/dict_to_schema/README.md new file mode 100644 index 0000000..847b69f --- /dev/null +++ b/examples/dict_to_schema/README.md @@ -0,0 +1,93 @@ +# Dict to Schema + +This example demonstrates how to automatically convert Python dictionary literals into Pydantic models. The codemod makes this process simple by handling all the tedious manual updates automatically. + +## How the Conversion Script Works + +The script (`run.py`) automates the entire conversion process in a few key steps: + +1. **Codebase Loading** + ```python + codebase = Codebase.from_repo("modal-labs/modal-client") + ``` + - Loads your codebase into Codegen's intelligent code analysis engine + - Provides a simple SDK for making codebase-wide changes + - Supports any Git repository as input + +2. **Dictionary Detection** + ```python + if "{" in global_var.source and "}" in global_var.source: + dict_content = global_var.value.source.strip("{}") + ``` + - Automatically identifies dictionary literals in your code + - Processes both global variables and class attributes + - Skips empty dictionaries to avoid unnecessary conversions + +3. **Schema Creation** + ```python + class_name = global_var.name.title() + "Schema" + model_def = f"""class {class_name}(BaseModel): + {dict_content.replace(",", "\n ")}""" + ``` + - Generates meaningful model names based on variable names + - Converts dictionary key-value pairs to class attributes + - Maintains proper Python indentation + +4. **Code Updates** + ```python + global_var.insert_before(model_def + "\n\n") + global_var.set_value(f"{class_name}(**{global_var.value.source})") + ``` + - Inserts new Pydantic models in appropriate locations + - Updates dictionary assignments to use the new models + - Automatically adds required Pydantic imports + + +## Common Conversion Patterns + +### Global Variables +```python +# Before +config = {"host": "localhost", "port": 8080} + +# After +class ConfigSchema(BaseModel): + host: str = "localhost" + port: int = 8080 + +config = ConfigSchema(**{"host": "localhost", "port": 8080}) +``` + +### Class Attributes +```python +# Before +class Service: + defaults = {"timeout": 30, "retries": 3} + +# After +class DefaultsSchema(BaseModel): + timeout: int = 30 + retries: int = 3 + +class Service: + defaults = DefaultsSchema(**{"timeout": 30, "retries": 3}) +``` + +## Running the Conversion + +```bash +# Install Codegen +pip install codegen + +# Run the conversion +python run.py +``` + +## Learn More + +- [Pydantic Documentation](https://docs.pydantic.dev/) +- [Codegen Documentation](https://docs.codegen.com) + +## Contributing + +Feel free to submit issues and enhancement requests! diff --git a/examples/dict_to_schema/run.py b/examples/dict_to_schema/run.py new file mode 100644 index 0000000..cfe505d --- /dev/null +++ b/examples/dict_to_schema/run.py @@ -0,0 +1,102 @@ +import codegen +from codegen import Codebase + + +@codegen.function("dict-to-pydantic-schema") +def run(codebase: Codebase): + """Convert dictionary literals to Pydantic models in a Python codebase. + + This codemod: + 1. Finds all dictionary literals in global variables and class attributes + 2. Creates corresponding Pydantic models + 3. Updates the assignments to use the new models + 4. Adds necessary Pydantic imports + """ + # Track statistics + files_modified = 0 + models_created = 0 + + # Iterate through all files in the codebase + for file in codebase.files: + needs_imports = False + file_modified = False + + # Look for dictionary assignments in global variables + for global_var in file.global_vars: + try: + if "{" in global_var.source and "}" in global_var.source: + dict_content = global_var.value.source.strip("{}") + if not dict_content.strip(): + continue + + # Convert dict to Pydantic model + class_name = global_var.name.title() + "Schema" + model_def = f"""class {class_name}(BaseModel): + {dict_content.replace(",", "\n ")}""" + + print(f"\nConverting '{global_var.name}' to schema") + print("\nOriginal code:") + print(global_var.source) + print("\nNew code:") + print(model_def) + print(f"{class_name}(**{global_var.value.source})") + print("-" * 50) + + # Insert model and update assignment + global_var.insert_before(model_def + "\n\n") + global_var.set_value(f"{class_name}(**{global_var.value.source})") + needs_imports = True + models_created += 1 + file_modified = True + except Exception as e: + print(f"Error processing global variable {global_var.name}: {str(e)}") + + # Look for dictionary assignments in class attributes + for cls in file.classes: + for attr in cls.attributes: + try: + if "{" in attr.source and "}" in attr.source: + dict_content = attr.value.source.strip("{}") + if not dict_content.strip(): + continue + + # Convert dict to Pydantic model + class_name = attr.name.title() + "Schema" + model_def = f"""class {class_name}(BaseModel): + {dict_content.replace(",", "\n ")}""" + + print(f"\nConverting'{attr.name}' to schema") + print("\nOriginal code:") + print(attr.source) + print("\nNew code:") + print(model_def) + print(f"{class_name}(**{attr.value.source})") + print("-" * 50) + + # Insert model and update attribute + cls.insert_before(model_def + "\n\n") + attr.set_value(f"{class_name}(**{attr.value.source})") + needs_imports = True + models_created += 1 + file_modified = True + except Exception as e: + print(f"Error processing attribute {attr.name} in class {cls.name}: {str(e)}") + + # Add imports if needed + if needs_imports: + file.add_import_from_import_string("from pydantic import BaseModel") + + if file_modified: + files_modified += 1 + + print("\nModification complete:") + print(f"Files modified: {files_modified}") + print(f"Schemas created: {models_created}") + + +if __name__ == "__main__": + print("Initializing codebase...") + codebase = Codebase.from_repo("modal-labs/modal-client") + + print("Running codemod...") + run(codebase) From ad1ca521ee0077b9dbaacc4d8bc19bd8115ff0c4 Mon Sep 17 00:00:00 2001 From: Vishal Shenoy Date: Tue, 28 Jan 2025 15:11:41 -0800 Subject: [PATCH 2/2] commit hash --- examples/dict_to_schema/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dict_to_schema/run.py b/examples/dict_to_schema/run.py index cfe505d..d8c633b 100644 --- a/examples/dict_to_schema/run.py +++ b/examples/dict_to_schema/run.py @@ -96,7 +96,7 @@ def run(codebase: Codebase): if __name__ == "__main__": print("Initializing codebase...") - codebase = Codebase.from_repo("modal-labs/modal-client") + codebase = Codebase.from_repo("modal-labs/modal-client", commit="81941c24897889a2ff2f627c693fa734967e693c") print("Running codemod...") run(codebase)