88from django .urls import reverse
99
1010from ..audits .models import Audit
11- from ..common .csv_export import CSVColumn , EqualityBodyCSVColumn , format_model_field
11+ from ..common .csv_export import (
12+ CSVColumn ,
13+ EqualityBodyCSVColumn ,
14+ ExportableClasses ,
15+ format_model_field ,
16+ )
1217from ..detailed .models import Contact as DetailedContact
1318from ..detailed .models import DetailedCase
1419from ..reports .models import Report
15- from ..simplified .csv_export import (
16- CONTACT_DETAILS_COLUMN_HEADER ,
17- SIMPLIFIED_EQUALITY_BODY_COLUMNS_FOR_EXPORT ,
18- format_simplified_contacts ,
19- )
20+ from ..simplified .csv_export import SIMPLIFIED_EQUALITY_BODY_COLUMNS_FOR_EXPORT
2021from ..simplified .models import CaseCompliance , CaseStatus
2122from ..simplified .models import Contact as SimplifiedContact
2223from ..simplified .models import SimplifiedCase
2324
2425DOWNLOAD_CASES_CHUNK_SIZE : int = 500
2526
27+ EqualityBodySourceClasses = (
28+ Audit | DetailedCase | CaseCompliance | Report | SimplifiedCase | None
29+ )
30+
2631
2732def populate_equality_body_columns (
2833 case : DetailedCase | SimplifiedCase ,
@@ -31,72 +36,64 @@ def populate_equality_body_columns(
3136 """
3237 Collect data for a case to export to the equality body
3338 """
34- contact_details : str = (
35- format_simplified_contacts (contacts = case .contacts )
36- if isinstance (case , SimplifiedCase )
37- else None
38- )
39- source_instances : dict = {
40- SimplifiedCase : case if isinstance (case , SimplifiedCase ) else None ,
41- Audit : case .audit if hasattr (case , "audit" ) else None ,
42- CaseCompliance : case .compliance if hasattr (case , "compliance" ) else None ,
43- Report : case .report if hasattr (case , "report" ) else None ,
44- }
39+ source_instances : dict [EqualityBodySourceClasses ] = {}
40+ if isinstance (case , DetailedCase ):
41+ source_instances [DetailedCase ] = case
42+ elif isinstance (case , SimplifiedCase ):
43+ source_instances [SimplifiedCase ] = case
44+ if hasattr (case , "audit" ):
45+ source_instances [Audit ] = case .audit
46+ if hasattr (case , "compliance" ):
47+ source_instances [CaseCompliance ] = case .compliance
48+ if hasattr (case , "report" ):
49+ source_instances [Report ] = case .report
50+
4551 columns : list [EqualityBodyCSVColumn ] = copy .deepcopy (column_definitions )
52+
4653 for column in columns :
47- source_instance : Audit | SimplifiedCase | CaseCompliance | Report | None = (
48- source_instances . get ( column .source_class )
54+ source_instance : EqualityBodySourceClasses = source_instances . get (
55+ column .source_class
4956 )
50- edit_url_instance : Audit | SimplifiedCase | CaseCompliance | Report | None = (
51- source_instances . get ( column .edit_url_class )
57+ edit_url_instance : EqualityBodySourceClasses = source_instances . get (
58+ column .edit_url_class
5259 )
53- if column .column_header == CONTACT_DETAILS_COLUMN_HEADER :
54- column .formatted_data = contact_details
55- else :
56- column .formatted_data = format_model_field (
57- source_instance = source_instance , column = column
58- )
60+
61+ column .formatted_data = format_model_field (
62+ source_instance = source_instance , column = column
63+ )
64+
5965 if column .edit_url_name is not None and edit_url_instance is not None :
6066 column .edit_url = reverse (
6167 column .edit_url_name , kwargs = {"pk" : edit_url_instance .id }
6268 )
6369 if column .edit_url_anchor :
6470 column .edit_url += f"#{ column .edit_url_anchor } "
71+
6572 return columns
6673
6774
6875def populate_csv_columns (
6976 case : DetailedCase | SimplifiedCase , column_definitions : list [CSVColumn ]
7077) -> list [CSVColumn ]:
71- """
72- Collect data for a case to export
73- """
74- source_instances : dict = {
75- SimplifiedCase : case if isinstance (case , SimplifiedCase ) else None ,
76- DetailedCase : case if isinstance (case , DetailedCase ) else None ,
77- CaseCompliance : case .compliance if hasattr (case , "compliance" ) else None ,
78- CaseStatus : case .status ,
79- SimplifiedContact : (
80- case .contact_set .filter (is_deleted = False ).first ()
81- if isinstance (case , SimplifiedCase )
82- else None
83- ),
84- DetailedContact : (
85- case .contact_set .filter (is_deleted = False ).first ()
86- if isinstance (case , DetailedCase )
87- else None
88- ),
89- }
78+ """Collect data for a case to export"""
79+ source_instances : dict [ExportableClasses ] = {}
80+ if isinstance (case , DetailedCase ):
81+ source_instances [DetailedCase ] = case
82+ source_instances [DetailedContact ] = case .contact_set .filter (
83+ is_deleted = False
84+ ).first ()
85+ elif isinstance (case , SimplifiedCase ):
86+ source_instances [SimplifiedCase ] = case
87+ source_instances [CaseStatus ] = case .status
88+ source_instances [SimplifiedContact ] = case .contact_set .filter (
89+ is_deleted = False
90+ ).first ()
91+ if hasattr (case , "compliance" ):
92+ source_instances [CaseCompliance ] = case .compliance
93+
9094 columns : list [CSVColumn ] = copy .deepcopy (column_definitions )
9195 for column in columns :
92- source_instance : (
93- DetailedCase
94- | SimplifiedCase
95- | CaseCompliance
96- | CaseStatus
97- | SimplifiedContact
98- | None
99- ) = source_instances .get (column .source_class )
96+ source_instance : ExportableClasses = source_instances .get (column .source_class )
10097 column .formatted_data = format_model_field (
10198 source_instance = source_instance , column = column
10299 )
0 commit comments