black fmt

Author: MatthewCaseres

pymort/XML.py

@@ -7,41 +7,49 @@
 
 # The following classes represent the xml elements in the XTbML file https://mort.soa.org/About.aspx
 
+
 @dataclass
 class AxisDef:
     """
     Corresponds to the XML element having the same name.
     """
+
     ScaleType: str
     AxisName: str
     MinScaleValue: int
     MaxScaleValue: int
     Increment: int
 
+
 @dataclass
 class MetaData:
     """
     Corresponds to the XML element having the same name.
     """
+
     ScalingFactor: float
     DataType: str
     Nation: str
     TableDescription: str
     AxisDefs: List[AxisDef]
 
+
 @dataclass
 class Table:
     """
     Corresponds to the XML element having the same name.
     """
+
     MetaData: MetaData
     Values: pd.DataFrame
 
+
 @dataclass
 class ContentClassification:
     """
     Corresponds to the XML element having the same name.
     """
+
     TableIdentity: int
     ProviderDomain: str
     ProviderName: str
@@ -52,11 +60,12 @@ class ContentClassification:
     Comments: str
     KeyWords: List[str]
 
+
 class MortXML:
     """A Python wrapper for XML mortality tables.
 
-    A tree of nested Python classes that represent an `XML standard <https://mort.soa.org/About.aspx/>`_ 
-    used by tables hosted at `mort.soa.org <https://mort.soa.org/>`_. 
+    A tree of nested Python classes that represent an `XML standard <https://mort.soa.org/About.aspx/>`_
+    used by tables hosted at `mort.soa.org <https://mort.soa.org/>`_.
 
     Corresponds to the `XTbML` root element in the XML file.
 
@@ -65,43 +74,50 @@ class MortXML:
         tables (List[Table]): Corresponds to Table XML elements.
 
     """
+
     def __init__(self, id: int):
         """
         Takes the id of a table and returns the PyXML object.
 
-        Args: 
+        Args:
             id (int): The id of the table to be loaded.
         """
         root = ET.fromstring(importlib.resources.read_text(data, f"t{id}.xml"))
-        self.ContentClassification: ContentClassification = createContentClassification(root.find('./ContentClassification'))
+        self.ContentClassification: ContentClassification = createContentClassification(
+            root.find("./ContentClassification")
+        )
         self.Tables: List[Table] = createTables(root)
 
+
 # The following functions turn XML elements into Python objects.
 
+
 def createAxisDef(axisDef: ET.Element) -> AxisDef:
-    '''
+    """
     Given an xml <AxisDef> element, return an AxisDef object
-    '''
+    """
     return AxisDef(
-        axisDef.find('./ScaleType').text,
-        axisDef.find('./AxisName').text,
-        int(axisDef.find('./MinScaleValue').text),
-        int(axisDef.find('./MaxScaleValue').text),
-        int(axisDef.find('./Increment').text)
+        axisDef.find("./ScaleType").text,
+        axisDef.find("./AxisName").text,
+        int(axisDef.find("./MinScaleValue").text),
+        int(axisDef.find("./MaxScaleValue").text),
+        int(axisDef.find("./Increment").text),
     )
 
+
 def createMetaData(metadata: ET.Element) -> MetaData:
-    '''
+    """
     Given an xml <MetaData> element, return a MetaData object
-    '''
+    """
     return MetaData(
-        float(metadata.find('./ScalingFactor').text),
-        metadata.find('./DataType').text,
-        metadata.find('./Nation').text,
-        metadata.find('./TableDescription').text,
-        [createAxisDef(axisDef) for axisDef in metadata.findall('./AxisDef')]
+        float(metadata.find("./ScalingFactor").text),
+        metadata.find("./DataType").text,
+        metadata.find("./Nation").text,
+        metadata.find("./TableDescription").text,
+        [createAxisDef(axisDef) for axisDef in metadata.findall("./AxisDef")],
     )
 
+
 def getAxisVals(axis: ET.Element) -> pd.DataFrame:
     """
     Take the bottom level `Axis` tag (having no "t" attribute).
@@ -115,14 +131,15 @@ def getAxisVals(axis: ET.Element) -> pd.DataFrame:
     if "t" in axis.attrib:
         rowT = int(axis.attrib["t"])
         colTs = ts
-        df = pd.DataFrame({"Age": [rowT]*len(vals), "Duration": colTs, "vals": vals})
+        df = pd.DataFrame({"Age": [rowT] * len(vals), "Duration": colTs, "vals": vals})
         df.set_index(["Age", "Duration"], inplace=True)
         return df
-    else: # one dimensional
+    else:  # one dimensional
         df = pd.DataFrame({"Age": ts, "vals": vals})
         df.set_index("Age", inplace=True)
         return df
 
+
 def createValues(table: ET.Element) -> pd.DataFrame:
     """
     Take the top level `Table` tag and return a list
@@ -132,31 +149,36 @@ def createValues(table: ET.Element) -> pd.DataFrame:
     for axis in axes:
         dfs.append(getAxisVals(axis))
     return pd.concat(dfs)
-    
+
+
 def createTable(table: ET.Element) -> Table:
-    '''
+    """
     Given an xml <Table> element, return a Table object
-    '''
-    metaData = createMetaData(table.find('./MetaData'))
+    """
+    metaData = createMetaData(table.find("./MetaData"))
     values = createValues(table)
     return Table(metaData, values)
 
+
 def createTables(root: ET.Element) -> List[Table]:
-    '''Given the root of an xml tree, return a list of Table objects'''
-    return [createTable(table) for table in root.findall('./Table')]
+    """Given the root of an xml tree, return a list of Table objects"""
+    return [createTable(table) for table in root.findall("./Table")]
+
 
-def createContentClassification(contentClassification: ET.Element) -> ContentClassification:
-    '''
+def createContentClassification(
+    contentClassification: ET.Element,
+) -> ContentClassification:
+    """
     Given an xml <ContentClassification> element, return a ContentClassification object
-    '''
+    """
     return ContentClassification(
-        int(contentClassification.find('./TableIdentity').text),
-        contentClassification.find('./ProviderDomain').text,
-        contentClassification.find('./ProviderName').text,
-        contentClassification.find('./TableReference').text,
-        contentClassification.find('./ContentType').text,
-        contentClassification.find('./TableName').text,
-        contentClassification.find('./TableDescription').text,
-        contentClassification.find('./Comments').text,
-        [keyword.text for keyword in contentClassification.findall('./KeyWord')]
-    )
+        int(contentClassification.find("./TableIdentity").text),
+        contentClassification.find("./ProviderDomain").text,
+        contentClassification.find("./ProviderName").text,
+        contentClassification.find("./TableReference").text,
+        contentClassification.find("./ContentType").text,
+        contentClassification.find("./TableName").text,
+        contentClassification.find("./TableDescription").text,
+        contentClassification.find("./Comments").text,
+        [keyword.text for keyword in contentClassification.findall("./KeyWord")],
+    )

pymort/__init__.py

@@ -2,5 +2,3 @@
 
 from .relational import Relational, getIdGroup
 from .XML import MortXML
-
-

pymort/archive_2022_May_04_093934/__init__.py

@@ -1 +1 @@
-"""These are the XML files retrieved at the date given in the directory name"""
+"""These are the XML files retrieved at the date given in the directory name"""

pymort/relational.py

@@ -6,6 +6,7 @@
 from typing import Tuple
 from itertools import groupby
 
+
 class Relational:
     """Rate tables with primary/foreign keys.
 
@@ -18,17 +19,26 @@ class Relational:
         metadata (pd.DataFrame): Composite primary key of `id`, `Age` is a MultiIndex. Column `vals` stores rates.
 
     """
+
     def __init__(self):
         # These pickled files are generated using `scripts/createRelational.py`.
-        self.metadata: pd.DataFrame = pickle.loads(importlib.resources.read_binary(pickled_tables, 'meta.pickle'))
-        self.select: pd.DataFrame = pickle.loads(importlib.resources.read_binary(pickled_tables, 'sel.pickle'))
-        self.ultimate: pd.DataFrame = pickle.loads(importlib.resources.read_binary(pickled_tables, 'ult.pickle'))
+        self.metadata: pd.DataFrame = pickle.loads(
+            importlib.resources.read_binary(pickled_tables, "meta.pickle")
+        )
+        self.select: pd.DataFrame = pickle.loads(
+            importlib.resources.read_binary(pickled_tables, "sel.pickle")
+        )
+        self.ultimate: pd.DataFrame = pickle.loads(
+            importlib.resources.read_binary(pickled_tables, "ult.pickle")
+        )
+
 
 @dataclass
 class IdGroup:
     """
     An IdGroup is a group of IDs that have the same study and grouping in the Relational().metadata pandas table
     """
+
     study: str
     grouping: str
     ids: Tuple[int]
@@ -45,9 +55,13 @@ def getIdGroup(targetId: int) -> IdGroup:
     meta.reset_index(inplace=True)
     # ensure that study/grouping groups are all consecutive
     meta.sort_values(by=["study", "grouping", "id"], inplace=True)
-    for k, g in groupby(zip(meta.study, meta.grouping, meta.id, meta.gender, meta.risk), key=lambda x: (x[0], x[1])):
+    for k, g in groupby(
+        zip(meta.study, meta.grouping, meta.id, meta.gender, meta.risk),
+        key=lambda x: (x[0], x[1]),
+    ):
         groupIds, genders, risks = zip(*[x[2:] for x in g])
         if targetId in groupIds:
             return IdGroup(k[0], k[1], groupIds, genders, risks)
-    raise KeyError("Your table identifier is not in the pandas table Relational().metadata")
-
+    raise KeyError(
+        "Your table identifier is not in the pandas table Relational().metadata"
+    )