Adding annotations - triangle.py

chainladder/core/base.py

@@ -1,6 +1,8 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
+from __future__ import annotations
+
 import pandas as pd
 from packaging import version
 
@@ -17,8 +19,15 @@
 from chainladder.core.io import TriangleIO
 from chainladder.core.common import Common
 from chainladder import options
-from chainladder.utils.utility_functions import num_to_nan, concat
 
+from typing import (
+    Optional,
+    TYPE_CHECKING
+)
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+    from pandas.core.interchange.dataframe_protocol import DataFrame as DataFrameXchg
 
 class TriangleBase(
     TriangleIO, TriangleDisplay, TriangleSlicer, TriangleDunders, TrianglePandas, Common
@@ -30,10 +39,22 @@ def shape(self):
         return self.values.shape
 
     @staticmethod
-    def _input_validation(data, index, columns, origin, development):
+    def _input_validation(
+            data: DataFrame,
+            index: str | list,
+            columns: str | list,
+            origin: str | list,
+            development: str | list
+    ) -> tuple[
+        None | list,
+        None | list,
+        None | list,
+        None | list
+    ]:
+
         """Validate/sanitize inputs"""
 
-        def str_to_list(arg):
+        def str_to_list(arg: str | list) -> None | list:
             if arg is None:
                 return
             if type(arg) in [str, pd.Period]:
@@ -219,7 +240,12 @@ def nan_triangle(self):
         return nan_triangle
 
     @staticmethod
-    def _to_datetime(data, fields, period_end=False, format=None):
+    def _to_datetime(
+            data: DataFrame,
+            fields: list,
+            period_end: bool = False,
+            format: Optional[str] = None
+    ):
         """For tabular form, this will take a set of data
         column(s) and return a single date array.  This function heavily
         relies on pandas, but does two additional things:
@@ -417,7 +443,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         else:
             raise NotImplementedError()
 
-    def _interchange_dataframe(self, data):
+    def _interchange_dataframe(self, data: DataFrameXchg) -> DataFrame:
         """
         Convert an object supporting the __dataframe__ protocol to a pandas DataFrame.
         Requires pandas version > 1.5.2.

chainladder/core/triangle.py

@@ -1,10 +1,10 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
+from __future__ import annotations
 
 import pandas as pd
 import numpy as np
-import copy
 import warnings
 from packaging import version
 from chainladder.core.base import TriangleBase
@@ -16,19 +16,31 @@
 
 try:
     import dask.bag as db
-except:
+except ImportError:
     db = None
 
+from typing import (
+    Optional,
+    TYPE_CHECKING
+)
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
+    from pandas.core.interchange.dataframe_protocol import DataFrame as DataFrameXchg
+
 
 class Triangle(TriangleBase):
     """
     The core data structure of the chainladder package
 
     Parameters
     ----------
-    data: DataFrame
-        A single dataframe that contains columns represeting all other
-        arguments to the Triangle constructor
+    data: DataFrame or DataFrameXchg
+        A single dataframe that contains columns representing all other
+        arguments to the Triangle constructor. If using pandas version > 1.5.2,
+        one may supply a DataFrame-like object (referred to as DataFrameXchg)
+        supporting the __dataframe__ protocol, which will then be converted to
+        a pandas DataFrame.
     origin: str or list
          A representation of the accident, reporting or more generally the
          origin period of the triangle that will map to the Origin dimension
@@ -108,35 +120,52 @@ class Triangle(TriangleBase):
 
     def __init__(
         self,
-        data=None,
-        origin=None,
-        development=None,
-        columns=None,
-        index=None,
-        origin_format=None,
-        development_format=None,
-        cumulative=None,
+        data: Optional[DataFrame | DataFrameXchg] = None,
+        origin: Optional[str | list] = None,
+        development: Optional[str | list] = None,
+        columns: Optional[str | list] = None,
+        index: Optional[str | list] = None,
+        origin_format: Optional[str] = None,
+        development_format: Optional[str] = None,
+        cumulative: Optional[bool] = None,
         array_backend=None,
         pattern=False,
-        trailing=True,
+        trailing: bool = True,
         *args,
         **kwargs
     ):
+
+        # If data are present, validate the dimensions.
         if data is None:
             return
         elif not isinstance(data, pd.DataFrame) and hasattr(data, "__dataframe__"):
             data = self._interchange_dataframe(data)
         index, columns, origin, development = self._input_validation(
-            data, index, columns, origin, development
+            data=data,
+            index=index,
+            columns=columns,
+            origin=origin,
+            development=development
         )
 
-        self.columns_label = columns
-        self.origin_label = origin
-
-        # Handle any ultimate vectors in triangles separately
-        data, ult = self._split_ult(data, index, columns, origin, development)
-        # Conform origins and developments to datetimes and determine lowest grains
-        origin_date = self._to_datetime(data, origin, format=origin_format).rename(
+        # Store dimension metadata.
+        self.columns_label: list = columns
+        self.origin_label: list = origin
+
+        # Handle any ultimate vectors in triangles separately.
+        data, ult = self._split_ult(
+            data=data,
+            index=index,
+            columns=columns,
+            origin=origin,
+            development=development
+        )
+        # Conform origins and developments to datetimes and determine the lowest grains
+        origin_date = self._to_datetime(
+            data=data,
+            fields=origin,
+            format=origin_format
+        ).rename(
             "__origin__"
         )
         self.origin_grain = self._get_grain(
@@ -192,8 +221,9 @@ def __init__(
         if cumulative is None:
             warnings.warn(
                 """
-            The cumulative property of your triangle is not set. This may result in
-            undesirable behavior. In a future release this will result in an error."""
+                The cumulative property of your triangle is not set. This may result in
+                undesirable behavior. In a future release this will result in an error.
+                """
             )
 
         self.is_cumulative = cumulative
@@ -278,7 +308,13 @@ def __init__(
             self.valuation_date = pd.Timestamp(options.ULT_VAL)
 
     @staticmethod
-    def _split_ult(data, index, columns, origin, development):
+    def _split_ult(
+            data: DataFrame,
+            index: list,
+            columns: list,
+            origin: list,
+            development: list
+    ) -> tuple[DataFrame, Triangle]:
         """Deal with triangles with ultimate values"""
         ult = None
         if (