Fixup type annotation

nsidc · Oct 2, 2024 · a04f64a · a04f64a
1 parent 2e31159
commit a04f64a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/usaon_benefit_tool/util/sankey.py b/usaon_benefit_tool/util/sankey.py
@@ -31,7 +31,7 @@ def permitted_source_link_types(node_type: NodeType) -> set[NodeType]:
     {
         "from": str,
         "to": str,
-        "weight": int,
+        "weight": int | float,
         "color": str,
         "id": NotRequired[int],
         "tooltipHTML": str,
@@ -180,7 +180,7 @@ def _node_ids_in_links(links: list[HighchartsSankeySeriesLink]) -> set[str]:
     return node_ids
 
 
-def _weight_for_criticality_rating(criticality_rating: int | None) -> float:
+def _weight_for_criticality_rating(criticality_rating: int | None) -> int | float:
     """If criticality rating is not set, return a very thin line."""
     if criticality_rating is None:
         return 0.01