Coverage for postrfp/shared/fetch/nodesq.py: 86%

1from typing import ( 

2 Any, 

3) 

4 

5from sqlalchemy import func, literal 

6from sqlalchemy.orm import ( 

7 Query, 

8 Session, 

9) 

10from postrfp.model import ( 

11 Section, 

12 QuestionInstance, 

13 QuestionDefinition, 

14) 

15from postrfp.model.questionnaire.b36 import from_b36 

16 

17 

18def light_nodes( 

19 session: Session, project_id: int, with_questions: bool = True 

20) -> Query: 

21 """ 

22 Returns a query of the project's questionnaire nodes. 

23 The query returns rows representing the sections and questions 

24 in the project. Each row has the following keys: 

25 - id: The ID of the section or question 

26 - title: The title of the section or question 

27 - b36_number: The b36_number of the section or question 

28 - type: The type of the node, either 'section' or 'question' 

29 - parent_id: The ID of the parent section of the section or question 

30 - depth: The depth of the section or question in the hierarchy 

31 

32 The query is ordered by the b36_number of the section or question. 

33 If the with_questions flag is set to True, the query will include 

34 the questions in the project. 

35 

36 This function avoids the expense of building ORM objects and is 

37 useful for building a nested structure of the project's questionnaire. 

38 """ 

39 

40 q = session.query( 

41 Section.id, 

42 Section.title, 

43 literal("section").label("type"), 

44 Section.parent_id, 

45 Section.b36_number.label("b36_number"), 

46 Section.position, 

47 (func.length(func.ifnull(Section.b36_number, "")) / 2).label("depth"), 

48 ).filter(Section.project_id == project_id) 

49 

50 if with_questions: 

51 aq = ( 

52 session.query( 

53 QuestionInstance.id, 

54 QuestionDefinition.title, 

55 literal("question").label("type"), 

56 QuestionInstance.section_id.label("parent_id"), 

57 QuestionInstance.b36_number.label("b36_number"), 

58 QuestionInstance.position, 

59 (func.length(QuestionInstance.b36_number) / 2).label("depth"), 

60 ) 

61 .join(QuestionDefinition) 

62 .filter(QuestionInstance.project_id == project_id) 

63 ) 

64 q = q.union(aq) 

65 

66 return q.order_by("b36_number") 

67 

68 

69def light_tree( 

70 session: Session, project_id: int, with_questions: bool = True 

71) -> dict[str, Any]: 

72 """ 

73 Returns a nested structure of the project's questionnaire. 

74 """ 

75 

76 node_lookup: dict[tuple[str, int], dict[str, Any]] = {} 

77 

78 for node in light_nodes(session, project_id, with_questions=with_questions): 

79 if node.type == "section": 

80 node_lookup[(node.type, node.id)] = { 

81 "id": node.id, 

82 "title": node.title, 

83 "number": from_b36(node.b36_number), 

84 "type": "section", 

85 "parent_id": node.parent_id, 

86 "depth": node.depth, 

87 "subsections": [], 

88 "questions": [], 

89 } 

90 else: 

91 node_lookup[(node.type, node.id)] = { 

92 "id": node.id, 

93 "title": node.title, 

94 "number": from_b36(node.b36_number), 

95 "type": "question", 

96 "parent_id": node.parent_id, 

97 "depth": node.depth, 

98 } 

99 

100 root_nodes: list[dict] = [] 

101 

102 for node in node_lookup.values(): 

103 parent_id = node["parent_id"] 

104 

105 if parent_id is None: 

106 root_nodes.append(node) 

107 continue 

108 

109 # Parent is always a section 

110 parent_key = ("section", parent_id) 

111 try: 

112 parent = node_lookup[parent_key] 

113 except KeyError: 

114 raise ValueError(f"Parent ID {parent_id} not found in node_lookup") 

115 

116 if node["type"] == "section": 

117 parent["subsections"].append(node) 

118 else: 

119 parent["questions"].append(node) 

120 

121 if len(root_nodes) == 1: 

122 return root_nodes[0] 

123 elif len(root_nodes) == 0: 

124 raise ValueError("No root found in the questionnaire") 

125 else: 

126 raise ValueError("Multiple roots found in the questionnaire")