Spaces:
Running
Running
| import json | |
| import ast | |
| import re | |
| def clean_and_refactor(code, function_name): | |
| # 1. First, UNWRAP anything that was double-wrapped or already wrapped | |
| # Remove any existing "class Solution:" lines and unindent | |
| lines = code.split('\n') | |
| unindented_lines = [] | |
| for line in lines: | |
| if line.strip() == "class Solution:": | |
| continue | |
| # Remove one level of indentation (4 spaces) if it was wrapped | |
| if line.startswith(" "): | |
| unindented_lines.append(line[4:]) | |
| else: | |
| unindented_lines.append(line) | |
| code = "\n".join(unindented_lines) | |
| # Remove redundant (self, self, ...) | |
| code = code.replace("(self, self, ", "(self, ") | |
| code = code.replace("(self, self)", "(self)") | |
| # 2. Apply the intended refactors (idempotent) | |
| code = code.replace("from tinygrad.tensor import Tensor", "import numpy as np") | |
| code = code.replace("tinygrad.tensor import Tensor", "numpy as np") | |
| # Replace Tensor( ... ) calls with np.array( ... ) - only if not torch. | |
| code = code.replace("Tensor(-1)", "-1") | |
| code = re.sub(r'(?<!torch\.)\bTensor\s*\(', 'np.array(', code) | |
| # Replace Tensor type hints with np.ndarray | |
| code = re.sub(r'[:]\s*(?<!torch\.)\bTensor\b', ': np.ndarray', code) | |
| code = re.sub(r'->\s*(?<!torch\.)\bTensor\b', '-> np.ndarray', code) | |
| # Catch-all for standalone Tensor | |
| code = re.sub(r'(?<!torch\.)\bTensor\b', 'np.ndarray', code) | |
| # 3. Re-wrap correctly | |
| lines = code.split('\n') | |
| imports = [] | |
| body = [] | |
| for line in lines: | |
| if line.strip().startswith(("import ", "from ")): | |
| imports.append(line) | |
| else: body.append(line) | |
| final_body = [] | |
| found_func = False | |
| for line in body: | |
| if line.strip().startswith("def "): | |
| # Ensure only ONE self is present | |
| if "(self, " not in line and "(self)" not in line: | |
| line = re.sub(r'def\s+(\w+)\s*\(', r'def \1(self, ', line) | |
| line = line.replace("self, )", "self)") | |
| found_func = True | |
| final_body.append(" " + line if (line.strip() or found_func) else line) | |
| return "\n".join(imports) + "\n\n" + "class Solution:\n" + "\n".join(final_body) | |
| def main(): | |
| with open("problems.json", "r") as f: | |
| problems = json.load(f) | |
| for p in problems: | |
| func_name = p.get("function_name", "unknown") | |
| if func_name == "unknown": | |
| match = re.search(r"def (\w+)\s*\(", p["starter_code"]) | |
| if match: func_name = match.group(1) | |
| p["function_name"] = func_name | |
| p["starter_code"] = clean_and_refactor(p["starter_code"], func_name) | |
| p["solution_code"] = clean_and_refactor(p["solution_code"], func_name) | |
| # Test cases are likely already fine (lists), but let's double check they aren't double-nested | |
| # Actually they looked okay in the previous view. | |
| with open("problems.json", "w") as f: | |
| json.dump(problems, f, indent=2) | |
| if __name__ == "__main__": | |
| main() | |