import wrapt


class ModelMatrix(wrapt.ObjectProxy):
    def __init__(self, data, spec=None):
        wrapt.ObjectProxy.__init__(self, data)
        self._self_model_spec = spec

    @property
    def model_spec(self):
        return self._self_model_spec

    def __repr__(self):
        return self.__wrapped__.__repr__()  # pragma: no cover

@@ -1,56 +1,618 @@
-import wrapt%0A%0A%0Aclass ModelMatrix(wrapt.ObjectProxy):
+from __future__ import annotations%0A%0Afrom typing import Any, Optional, TYPE_CHECKING%0A%0Aimport wrapt%0A%0Aif TYPE_CHECKING:  # pragma: no cover%0A    from .model_spec import ModelSpec%0A%0A%0Aclass ModelMatrix(wrapt.ObjectProxy):%0A    %22%22%22%0A    A wrapper around arbitrary model matrix output representations.%0A%0A    This wrapper allows for %60isinstance(..., ModelMatrix)%60 checks, and allows%0A    one to access the %60ModelSpec%60 instance associated with its creation using%0A    %60%3Cmodel_matrix%3E.model_spec%60. All other instance attributes and methods of%0A    the wrapped object are directly accessible as if the object were unwrapped.%0A    %22%22%22%0A
 %0A   
@@ -635,19 +635,49 @@
 lf, 
-data, spec=
+matrix: Any, spec: Optional%5BModelSpec%5D = 
 None
@@ -724,12 +724,14 @@
 lf, 
-data
+matrix
 )%0A  
@@ -800,25 +800,408 @@
 l_spec(self)
-:
+ -%3E Optional%5BModelSpec%5D:%0A        %22%22%22%0A        The %60ModelSpec%60 instance associated with the creation of this%0A        %60ModelMatrix%60 instance.%0A%0A        This %60ModelSpec%60 instance can be used to create other %60ModelMatrix%60s%0A        that respect all the choices (including feature selection and encoding)%0A        that were made in the construction of this %60ModelMatrix%60 instance.%0A%0A        %22%22%22
 %0A        ret

Document and annotate `ModelMatrix`.