From b2b95b347e4b1c1545d0d64527be069032f0bcc2 Mon Sep 17 00:00:00 2001
From: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com>
Date: Fri, 24 Nov 2023 06:44:32 +1300
Subject: [PATCH] Rewrite list-map-zip for mypy

Guido, and hence Python, do not favor lamba expressions.  Among other things, lambdas can confuse mypy and there's no great syntactic sugar to fix this.

We fix this two ways.  In the simple case, we can provide a type signature using typing.cast.  The a more complex case, we rewrite a nested comprehension that was written to create three lists as a simple loop to create those three lists with proper type signatures.  And in the process, simplify some old code that is no longer relevant.

Signed-off-by: Michael Tiemann <72577720+MichaelTiemannOSC@users.noreply.github.com>
---
 src/ITR/data/base_providers.py | 67 ++++++++++++++--------------------
 1 file changed, 27 insertions(+), 40 deletions(-)

diff --git a/src/ITR/data/base_providers.py b/src/ITR/data/base_providers.py
index e306daa9..2c05da75 100644
--- a/src/ITR/data/base_providers.py
+++ b/src/ITR/data/base_providers.py
@@ -2,7 +2,7 @@
 import warnings  # needed until quantile behaves better with Pint quantities in arrays
 from functools import partial, reduce
 from operator import add
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Callable, Dict, List, Optional, Type, cast
 
 import numpy as np
 import pandas as pd
@@ -889,45 +889,32 @@ def get_company_projected_trajectories(self, company_ids: List[str], year=None)
         :param year: values for a specific year, or all years if None
         :return: A pandas DataFrame with projected intensity trajectories per company, indexed by company_id and scope
         """
-        company_ids, scopes, projections = list(
-            map(
-                list,
-                zip(
-                    *[
-                        (
-                            c.company_id,
-                            EScope[scope_name],
-                            c.projected_intensities[scope_name].projections,
-                        )
-                        # FIXME: we should make _companies a dict so we can look things up rather than searching every time!
-                        for c in self._companies
-                        for scope_name in EScope.get_scopes()
-                        if c.company_id in company_ids
-                        if c.projected_intensities[scope_name]
-                    ]
-                ),
-            )
-        )
-        if projections:
-            index = pd.MultiIndex.from_tuples(zip(company_ids, scopes), names=["company_id", "scope"])
-            if year is not None:
-                if isinstance(projections[0], ICompanyEIProjectionsScopes):
-                    values = [yvp.value for pt in projections for yvp in pt if yvp.year == year]
-                else:
-                    values = list(map(lambda x: x[year].squeeze(), projections))
-                with warnings.catch_warnings():
-                    warnings.simplefilter("ignore")
-                    # pint units don't like columns of heterogeneous data...tough!
-                    return pd.Series(data=values, index=index, name=year)
-            else:
-                if isinstance(projections[0], ICompanyEIProjectionsScopes):
-                    values = [{yvp.year: yvp.value for yvp in pt} for pt in projections]
-                else:
-                    values = projections
-                with warnings.catch_warnings():
-                    warnings.simplefilter("ignore")
-                    return pd.DataFrame(data=values, index=index)
-        return pd.DataFrame()
+        c_ids: List[str] = []
+        scopes: List[EScope] = []
+        projections: List[DF_ICompanyEIProjections] = []
+
+        for c in self._companies:
+            if c.company_id in company_ids:
+                for scope_name in EScope.get_scopes():
+                    if c.projected_intensities[scope_name]:
+                        c_ids.append(c.company_id)
+                        scopes.append(EScope[scope_name])
+                        projections.append(c.projected_intensities[scope_name].projections)
+
+        if len(projections) == 0:
+            return pd.DataFrame()
+        index = pd.MultiIndex.from_tuples(zip(c_ids, scopes), names=["company_id", "scope"])
+        if year is not None:
+            values = list(map(cast(Callable[[pd.Series], Any], lambda x: x[year].squeeze()), projections))
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                # pint units don't like columns of heterogeneous data...tough!
+                return pd.Series(data=values, index=index, name=year)
+        else:
+            values = projections
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                return pd.DataFrame(data=values, index=index)
 
     def get_company_projected_targets(self, company_ids: List[str], year=None) -> pd.DataFrame:
         """