from intake.source.base import DataSource, Schema
import joblib
import fsspec
import sklearn
import re

from . import __version__

class SklearnModelSource(DataSource):
    container = 'python'
    name = 'sklearn'
    version = __version__
    partition_access = False

    def __init__(self, urlpath, storage_options=None, metadata=None):
        """
        Parameters
        ----------

        urlpath: str, location of model pkl file
        Either the absolute or relative path to the file or URL to be
        opened. Some examples:
          - ``{{ CATALOG_DIR }}/models/model.pkl``
          - ``s3://some-bucket/models/model.pkl``
        """

        self._urlpath = urlpath
        self._storage_options = storage_options or {}

        super().__init__(metadata=metadata)


    def _load(self):
        with fsspec.open(self._urlpath, mode='rb', **self._storage_options) as f:
            return f.read()


    def _get_schema(self):
        as_binary = self._load()
        s = re.search(b'_sklearn_versionq(.*\x00)((\d+\.)?(\d+\.)?(\*|\d+))q', as_binary)
        if s:
            sklearn_version = s.group(2).decode()
        else:
            sklearn_version = None

        self._schema = Schema(
            npartitions=1,
            extra_metadata={
                'sklearn_version':sklearn_version
            }
        )
        return self._schema


    def read(self):
        self._load_metadata()

        if not self.metadata['sklearn_version'] == sklearn.__version__:
            msg = ('The model was created with Scikit-Learn version {} '
                   'but version {} has been installed in your current environment.'
                  ).format(self.metadata['sklearn_version'], sklearn.__version__)
            raise RuntimeError(msg)


        with fsspec.open(self._urlpath, **self._storage_options) as f:
            return joblib.load(f)


[((40, 12, 40, 89), 're.search', 're.search', ({(40, 22, 40, 77): "b'_sklearn_versionq(.*\\x00)((\\\\d+\\\\.)?(\\\\d+\\\\.)?(\\\\*|\\\\d+))q'", (40, 79, 40, 88): 'as_binary'}, {}), "(b'_sklearn_versionq(.*\\x00)((\\\\d+\\\\.)?(\\\\d+\\\\.)?(\\\\*|\\\\d+))q',\n    as_binary)", False, 'import re\n'), ((46, 23, 51, 9), 'intake.source.base.Schema', 'Schema', (), '', False, 'from intake.source.base import DataSource, Schema\n'), ((34, 13, 34, 75), 'fsspec.open', 'fsspec.open', (), '', False, 'import fsspec\n'), ((65, 13, 65, 64), 'fsspec.open', 'fsspec.open', ({(65, 25, 65, 38): 'self._urlpath'}, {}), '(self._urlpath, **self._storage_options)', False, 'import fsspec\n'), ((66, 19, 66, 33), 'joblib.load', 'joblib.load', ({(66, 31, 66, 32): 'f'}, {}), '(f)', False, 'import joblib\n')]