def __init__(
self,
task_name: str,
electorate_layer: QgsVectorLayer,
meshblock_layer: QgsVectorLayer, # pylint: disable=too-many-locals, too-many-statements
meshblock_number_field_name: str,
scenario_registry: ScenarioRegistry,
scenario,
task: Optional[str] = None):
"""
Constructor for ScenarioSwitchTask
:param task_name: user-visible, translated name for task
:param electorate_layer: electorate layer
:param meshblock_layer: meshblock layer
:param meshblock_number_field_name: name of meshblock number field
:param scenario_registry: scenario registry
:param scenario: target scenario id to switch to
:param task: current redistricting task
"""
super().__init__(task_name)
self.scenario = scenario
self.electorate_layer = electorate_layer
self.task = task
self.type_idx = electorate_layer.fields().lookupField('type')
assert self.type_idx >= 0
self.scenario_id_idx = electorate_layer.fields().lookupField(
'scenario_id')
assert self.scenario_id_idx >= 0
self.estimated_pop_idx = electorate_layer.fields().lookupField(
'estimated_pop')
assert self.estimated_pop_idx >= 0
self.mb_number_idx = scenario_registry.meshblock_electorate_layer.fields(
).lookupField('meshblock_number')
assert self.mb_number_idx >= 0
self.mb_off_pop_m_idx = meshblock_layer.fields().lookupField(
'offline_pop_m')
assert self.mb_off_pop_m_idx >= 0
self.mb_off_pop_ni_idx = meshblock_layer.fields().lookupField(
'offline_pop_gn')
assert self.mb_off_pop_ni_idx >= 0
self.mb_off_pop_si_idx = meshblock_layer.fields().lookupField(
'offline_pop_gs')
assert self.mb_off_pop_si_idx >= 0
self.mb_offshore_idx = meshblock_layer.fields().lookupField('offshore')
assert self.mb_offshore_idx >= 0
self.invalid_reason_idx = self.electorate_layer.fields().lookupField(
'invalid_reason')
assert self.invalid_reason_idx >= 0
self.invalid_idx = self.electorate_layer.fields().lookupField(
'invalid')
assert self.invalid_idx >= 0
electorate_id_idx = electorate_layer.fields().lookupField(
'electorate_id')
assert electorate_id_idx >= 0
self.code_idx = electorate_layer.fields().lookupField('code')
assert self.code_idx >= 0
self.name_idx = electorate_layer.fields().lookupField('name')
assert self.name_idx >= 0
self.meshblock_number_idx = meshblock_layer.fields().lookupField(
meshblock_number_field_name)
assert self.meshblock_number_idx >= 0
# do a bit of preparatory processing on the main thread for safety
# dict of meshblock number to feature
meshblocks = {}
for m in meshblock_layer.getFeatures():
meshblocks[int(m[self.meshblock_number_idx])] = m
# dict of electorates to process (by id)
self.electorates_to_process = {}
request = QgsFeatureRequest().setFlags(QgsFeatureRequest.NoGeometry)
request.setSubsetOfAttributes(
[electorate_id_idx, self.type_idx, self.code_idx, self.name_idx])
for electorate in electorate_layer.getFeatures(request):
# get meshblocks for this electorate in the target scenario
electorate_id = electorate[electorate_id_idx]
electorate_type = electorate[self.type_idx]
electorate_code = electorate[self.code_idx]
electorate_name = electorate[self.name_idx]
if self.task and electorate_type != self.task:
continue
electorate_meshblocks = scenario_registry.electorate_meshblocks(
electorate_id=electorate_id,
electorate_type=electorate_type,
scenario_id=scenario)
assigned_meshblock_numbers = [
m[self.mb_number_idx] for m in electorate_meshblocks
]
matching_meshblocks = [
meshblocks[m] for m in assigned_meshblock_numbers
]
offshore_meshblocks = [
m for m in matching_meshblocks if m[self.mb_offshore_idx]
]
non_offshore_meshblocks = [
m for m in matching_meshblocks if not m[self.mb_offshore_idx]
]
self.electorates_to_process[electorate_id] = {
self.ELECTORATE_FEATURE_ID: electorate.id(),
self.ELECTORATE_TYPE: electorate_type,
self.ELECTORATE_CODE: electorate_code,
self.ELECTORATE_NAME: electorate_name,
self.MESHBLOCKS: matching_meshblocks,
self.OFFSHORE_MESHBLOCKS: offshore_meshblocks,
self.NON_OFFSHORE_MESHBLOCKS: non_offshore_meshblocks
}
self.setDependentLayers([electorate_layer])
def testElectorateMeshblocks(self):
"""
Test retrieving meshblocks belong to an electorate for a scenario
"""
layer = make_scenario_layer()
mb_electorate_layer = make_meshblock_electorate_layer()
reg = ScenarioRegistry(source_layer=layer,
id_field='id',
name_field='name',
meshblock_electorate_layer=mb_electorate_layer)
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='c', electorate_type='GN', scenario_id=1)
]
self.assertEqual(res, [0])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='d', electorate_type='GN', scenario_id=1)
]
self.assertEqual(res, [1])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='a', electorate_type='GN', scenario_id=1)
]
self.assertEqual(res, [])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='z', electorate_type='GS', scenario_id=1)
]
self.assertEqual(res, [0])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='zz', electorate_type='GS', scenario_id=1)
]
self.assertEqual(res, [1])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='x', electorate_type='GS', scenario_id=1)
]
self.assertEqual(res, [])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='a', electorate_type='GN', scenario_id=2)
]
self.assertEqual(res, [0])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='b', electorate_type='GN', scenario_id=2)
]
self.assertEqual(res, [1])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='c', electorate_type='GN', scenario_id=2)
]
self.assertEqual(res, [])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='x', electorate_type='GS', scenario_id=2)
]
self.assertEqual(res, [0])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='y', electorate_type='GS', scenario_id=2)
]
self.assertEqual(res, [1])
res = [
f['meshblock_number'] for f in reg.electorate_meshblocks(
electorate_id='z', electorate_type='GS', scenario_id=2)
]
self.assertEqual(res, [])