Coverage for class_generator/parsers/user_code_parser.py: 100%

1"""Parser for extracting user-added code from generated files.""" 

2 

3import ast 

4from typing import Any 

5 

6 

7def parse_user_code_from_file(file_path: str) -> tuple[str, str]: 

8 """ 

9 Extract user-added code and imports from a generated file. 

10  

11 Args: 

12 file_path: Path to the generated file 

13  

14 Returns: 

15 Tuple of (user_code, user_imports) 

16 """ 

17 with open(file_path) as fd: 

18 data = fd.read() 

19 

20 end_of_generated_code_line = " # End of generated code" 

21 user_code: str = "" 

22 user_imports: str = "" 

23 

24 if end_of_generated_code_line in data: 

25 _end_of_generated_code_index = data.index(end_of_generated_code_line) 

26 user_code = data[_end_of_generated_code_index + len(end_of_generated_code_line) :] 

27 

28 tree = ast.parse(data) 

29 imports = [imp for imp in tree.body if isinstance(imp, ast.Import) or isinstance(imp, ast.ImportFrom)] 

30 splited_data = data.splitlines() 

31 

32 # Standard imports that are always generated by the template 

33 template_imports = { 

34 "from typing import Any", 

35 "from ocp_resources.resource import NamespacedResource", 

36 "from ocp_resources.resource import Resource", 

37 "from ocp_resources.resource import NamespacedResource, MissingRequiredArgumentError", 

38 "from ocp_resources.resource import Resource, MissingRequiredArgumentError", 

39 } 

40 

41 for _import in imports: 

42 end_import_no = _import.end_lineno 

43 import_lines = [] 

44 

45 if end_import_no and _import.lineno != end_import_no: 

46 for num in range(_import.lineno - 1, end_import_no): 

47 import_lines.append(splited_data[num]) 

48 else: 

49 import_lines.append(splited_data[_import.lineno - 1]) 

50 

51 import_str = "\n".join(import_lines).strip() 

52 

53 # Only include imports that are not in the template 

54 if import_str not in template_imports: 

55 user_imports += f"{import_str}\n" 

56 

57 return user_code, user_imports