base

class BaseReader(ABC):
    """Abstract reader to parse various inputs into neurons.

    Any subclass should implement at least one of `read_buffer` or
    `read_dataframe`. Entry methods such as `read_any` will pass
    and parse an input through to the appropriate method.

    Parameters
    ----------
    fmt :           str
                    A string describing how to parse filenames into neuron
                    properties. For example '{id}.swc'.
    file_ext :      str | tuple
                    The file extension(s) to look for when searching folders.
                    For example '.swc'. Alternatively, you can re-implement
                    the `is_valid_file` method for more complex filters. That
                    method needs to be able to deal with: Path objects, ZipInfo
                    objects and strings.
    name_fallback : str
                    Fallback for name when reading from e.g. string.
    attrs :         dict
                    Additional attributes to use when creating the neuron.
                    Will be overwritten by later additions (e.g. from `fmt`).
    ignore_hidden : bool
                    Whether to ignore files that start with "._".
    errors :        "raise" | "log" | "ignore"
                    What to do when an error is encountered:
                     - "raise" (default) will raise an error
                     - "log" will log a warning and return `None`
                     - "ignore" will return `None`
                    Applies only to errors in parsing file contents into neurons
                    not to errors in reading files, archives, URLs, etc.

    """

    def __init__(
        self,
        fmt: str,
        file_ext: str,
        name_fallback: str = "NA",
        read_binary: bool = False,
        attrs: Optional[Dict[str, Any]] = None,
        ignore_hidden=True,
        errors="raise"
    ):
        self.attrs = attrs
        self.fmt = fmt
        self.file_ext = file_ext
        self.name_fallback = name_fallback
        self.read_binary = read_binary
        self.ignore_hidden = ignore_hidden
        self.errors = errors

        assert errors in ("raise", "log", "ignore")

    @property
    def file_ext(self):
        return self._file_ext

    @file_ext.setter
    def file_ext(self, value):
        """Makes sure file_ext is always a tuple."""
        if isinstance(value, str):
            value = (value,)

        if any((ext.startswith("*") for ext in value)):
            raise ValueError('File extension must be ".ext", not "*.ext"')

        self._file_ext = value

    def format_output(self, x):
        """Format output into NeuronList.

        Replace this method if output is not (always) a NeuronList.
        See for example NrrdReader in nrrd_io.py.
        """
        if not x:
            return core.NeuronList([])
        else:
            return core.NeuronList([n for n in x if n])

    def files_in_dir(
        self, dpath: Path, include_subdirs: bool = DEFAULT_INCLUDE_SUBDIRS
    ) -> Iterable[Path]:
        """List files to read in directory."""
        if not isinstance(dpath, Path):
            dpath = Path(dpath)

        if "*" in str(dpath):
            pattern = str(dpath.name)
            dpath = dpath.parent
        else:
            pattern = "*"

        dpath = dpath.expanduser()
        if include_subdirs:
            pattern = os.path.join("**", pattern)

        yield from (f for f in dpath.glob(pattern) if self.is_valid_file(f))

    def is_valid_file(self, file):
        """Return true if file should be considered for reading."""
        if isinstance(file, ZipInfo):
            file = file.filename
        elif isinstance(file, tarfile.TarInfo):
            file = file.name
        elif isinstance(file, Path):
            file = file.name

        if self.ignore_hidden and str(file).startswith("._"):
            return False

        for ext in self.file_ext:
            if str(file).endswith(ext):
                return True
        return False

    def _make_attributes(
        self, *dicts: Optional[Dict[str, Any]], **kwargs
    ) -> Dict[str, Any]:
        """Combine attributes with a timestamp and those defined on the object.

        Later additions take precedence:

        - created_at (now)
        - object-defined attributes
        - additional dicts given as *args
        - additional attributes given as **kwargs

        Returns
        -------
        Dict[str, Any]
            Arbitrary string-keyed attributes.
        """
        return merge_dicts(
            dict(created_at=str(datetime.datetime.now())),
            self.attrs,
            *dicts,
            **kwargs,
        )

    def read_file_path(
        self, fpath: os.PathLike, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Read single file from path into a neuron.

        Parameters
        ----------
        fpath :     str | os.PathLike
                    Path to files.
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        p = Path(fpath)
        with open(p, "rb" if self.read_binary else "r") as f:
            props = self.parse_filename(f.name)
            props["origin"] = str(p)
            return self.read_buffer(f, attrs=merge_dicts(props, attrs))

    def read_from_zip(
        self,
        files: Union[str, List[str]],
        zippath: os.PathLike,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read given files from a zip into a NeuronList.

        Typically not used directly but via `read_zip()` dispatcher.

        Parameters
        ----------
        files :     zipfile.ZipInfo | list thereof
                    Files inside the ZIP file to read.
        zippath :   str | os.PathLike
                    Path to zip file.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        p = Path(zippath)
        files = utils.make_iterable(files)

        neurons = []
        with ZipFile(p, "r") as zip:
            for file in files:
                # Note the `file` is of type zipfile.ZipInfo here
                props = self.parse_filename(file.orig_filename)
                props["origin"] = str(p)
                try:
                    n = self.read_bytes(zip.read(file), attrs=merge_dicts(props, attrs))
                    neurons.append(n)
                except BaseException:
                    if on_error == "ignore":
                        logger.warning(f'Failed to read "{file.filename}" from zip.')
                    else:
                        raise

        return self.format_output(neurons)

    def read_zip(
        self,
        fpath: os.PathLike,
        parallel="auto",
        limit: Optional[int] = None,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read files from a zip into a NeuronList.

        This is a dispatcher for `.read_from_zip`.

        Parameters
        ----------
        fpath :     str | os.PathLike
                    Path to zip file.
        limit :     int, optional
                    Limit the number of files read from this directory.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        fpath = Path(fpath).expanduser()
        read_fn = partial(
            self.read_from_zip, zippath=fpath, attrs=attrs, on_error=on_error
        )
        neurons = parallel_read_archive(
            read_fn=read_fn,
            fpath=fpath,
            file_ext=self.is_valid_file,
            limit=limit,
            parallel=parallel,
        )
        return self.format_output(neurons)

    def read_from_tar(
        self,
        files: Union[str, List[str]],
        tarpath: os.PathLike,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read given files from a tar into a NeuronList.

        Typically not used directly but via `read_tar()` dispatcher.

        Parameters
        ----------
        files :     tarfile.TarInfo | list thereof
                    Files inside the tar file to read.
        tarpath :   str | os.PathLike
                    Path to tar file.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        p = Path(tarpath)
        files = utils.make_iterable(files)

        neurons = []
        with tarfile.open(p, "r") as tf:
            for file in files:
                # Note the `file` is of type tarfile.TarInfo here
                props = self.parse_filename(file.name.split("/")[-1])
                props["origin"] = str(p)
                try:
                    n = self.read_bytes(
                        tf.extractfile(file).read(), attrs=merge_dicts(props, attrs)
                    )
                    neurons.append(n)
                except BaseException:
                    if on_error == "ignore":
                        logger.warning(f'Failed to read "{file.filename}" from tar.')
                    else:
                        raise

        return self.format_output(neurons)

    def read_tar(
        self,
        fpath: os.PathLike,
        parallel="auto",
        limit: Optional[int] = None,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read files from a tar archive into a NeuronList.

        This is a dispatcher for `.read_from_tar`.

        Parameters
        ----------
        fpath :     str | os.PathLike
                    Path to tar file.
        limit :     int, optional
                    Limit the number of files read from this directory.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        fpath = Path(fpath).expanduser()
        read_fn = partial(
            self.read_from_tar, tarpath=fpath, attrs=attrs, on_error=on_error
        )
        neurons = parallel_read_archive(
            read_fn=read_fn,
            fpath=fpath,
            file_ext=self.is_valid_file,
            limit=limit,
            parallel=parallel,
        )
        return self.format_output(neurons)

    def read_ftp(
        self,
        url,
        parallel="auto",
        limit: Optional[int] = None,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read files from an FTP server.

        This is a dispatcher for `.read_from_ftp`.

        Parameters
        ----------
        url :       str
                    Can be the path to a single file or a directory.
        limit :     int, optional
                    Limit the number of files read from this directory.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        # Remove the ftp:// prefix
        url = url.replace("ftp://", "")

        # Split into server and path
        server, path = url.split("/", 1)

        # Check if server contains a port
        if ":" in server:
            server, port = server.split(":")
            port = int(port)
        else:
            port = 21  # default port

        read_fn = partial(self.read_from_ftp, attrs=attrs, on_error=on_error)
        neurons = parallel_read_ftp(
            read_fn=read_fn,
            server=server,
            port=port,
            path=path,
            file_ext=self.is_valid_file,
            limit=limit,
            parallel=parallel,
        )
        return self.format_output(neurons)

    def read_from_ftp(
        self,
        files: Union[str, List[str]],
        ftp: FTP,
        attrs: Optional[Dict[str, Any]] = None,
        on_error: Union[Literal["ignore", Literal["raise"]]] = "ignore",
    ) -> "core.NeuronList":
        """Read given files from an FTP server into a NeuronList.

        Typically not used directly but via `read_ftp()` dispatcher.

        Parameters
        ----------
        files :     tarfile.TarInfo | list thereof
                    Files inside the tar file to read.
        ftp :       ftplib.FTP | "GLOBAL"
                    The FTP client. This should already be connected, logged in
                    and in the correct directory. If "GLOBAL", we will look for a
                    `_FTP` global variable.
        attrs :     dict or None
                    Arbitrary attributes to include in the TreeNeuron.
        on_error :  'ignore' | 'raise'
                    What do do when error is encountered.

        Returns
        -------
        core.NeuronList

        """
        # When reading in parallel, we expect there to be a global FTP connection
        # that was initialized once for each worker process.
        if ftp == "GLOBAL":
            if "_FTP" not in globals():
                raise ValueError("No global FTP connection found.")
            ftp = _FTP

        files = utils.make_iterable(files)

        neurons = []
        for file in files:
            # Read the file into a bytes
            with io.BytesIO() as f:
                ftp.retrbinary("RETR " + file, f.write)
                f.seek(0)
                props = self.parse_filename(file)
                props["origin"] = f"{ftp.host}:{ftp.port}{ftp.pwd()}/{file}"
                try:
                    n = self.read_buffer(f, attrs=merge_dicts(props, attrs))
                    neurons.append(n)
                except BaseException:
                    if on_error == "ignore":
                        logger.warning(f'Failed to read "{file}" from FTP.')
                    else:
                        raise

        return self.format_output(neurons)

    def read_directory(
        self,
        path: os.PathLike,
        include_subdirs=DEFAULT_INCLUDE_SUBDIRS,
        parallel="auto",
        limit: Optional[int] = None,
        attrs: Optional[Dict[str, Any]] = None,
    ) -> "core.NeuronList":
        """Read directory of files into a NeuronList.

        Parameters
        ----------
        fpath :             str | os.PathLike
                            Path to directory containing files.
        include_subdirs :   bool, optional
                            Whether to descend into subdirectories, default False.
        parallel :          str | bool | "auto"
        limit :             int, optional
                            Limit the number of files read from this directory.
        attrs :             dict or None
                            Arbitrary attributes to include in the TreeNeurons
                            of the NeuronList

        Returns
        -------
        core.NeuronList

        """
        files = list(self.files_in_dir(Path(path), include_subdirs))

        if isinstance(limit, int):
            files = files[:limit]
        elif isinstance(limit, list):
            files = [f for f in files if f in limit]
        elif isinstance(limit, slice):
            files = files[limit]
        elif isinstance(limit, str):
            # Check if limit is a regex
            if rgx.search(limit):
                files = [f for f in files if re.search(limit, str(f.name))]
            else:
                files = [f for f in files if limit in str(f)]

        read_fn = partial(self.read_file_path, attrs=attrs)
        neurons = parallel_read(read_fn, files, parallel)
        return self.format_output(neurons)

    def read_url(
        self, url: str, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Read file from URL into a neuron.

        Parameters
        ----------
        url :       str
                    URL to file.
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        # Note: originally, we used stream=True and passed `r.raw` to the
        # read_buffer function but that caused issue when there was more
        # than one chunk which would require us to concatenate the chunks
        # `via r.raw.iter_content()`.
        # Instead, we will simply read the whole content, wrap it in a BytesIO
        # and pass that to the read_buffer function. This is not ideal as it
        # will load the whole file into memory while the streaming solution
        # may have raised an exception earlier if the file was corrupted or
        # the wrong format.
        with requests.get(url, stream=False) as r:
            r.raise_for_status()
            props = self.parse_filename(url.split("/")[-1])
            props["origin"] = url
            return self.read_buffer(io.BytesIO(r.content), attrs=merge_dicts(props, attrs))

    def read_string(
        self, s: str, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Read single string into a Neuron.

        Parameters
        ----------
        s :         str
                    String.
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        sio = io.StringIO(s)
        return self.read_buffer(
            sio, attrs=merge_dicts({"name": self.name_fallback, "origin": "string"}, attrs)
        )

    def read_bytes(
        self, s: str, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Read bytes into a Neuron.

        Parameters
        ----------
        s :         bytes
                    Bytes.
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        sio = io.BytesIO(s)
        return self.read_buffer(
            sio, attrs=merge_dicts({"name": self.name_fallback, "origin": "string"}, attrs)
        )

    @handle_errors
    def read_dataframe(
        self, nodes: pd.DataFrame, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Convert a DataFrame into a neuron.

        Parameters
        ----------
        nodes :     pandas.DataFrame
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        raise NotImplementedError(
            "Reading DataFrames not implemented for " f"{type(self)}"
        )

    @handle_errors
    def read_buffer(
        self, f: IO, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Read buffer into a single neuron.



        Parameters
        ----------
        f :         IO
                    Readable buffer.
        attrs :     dict | None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.NeuronObject
        """
        raise NotImplementedError(
            "Reading from buffer not implemented for " f"{type(self)}"
        )

    def read_any_single(
        self, obj, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.BaseNeuron":
        """Attempt to convert an arbitrary object into a neuron.

        Parameters
        ----------
        obj :       typing.IO | pandas.DataFrame | str | os.PathLike
                    Readable buffer, dataframe, path or URL to single file,
                    or parsable string.
        attrs :     dict or None
                    Arbitrary attributes to include in the neuron.

        Returns
        -------
        core.BaseNeuron
        """
        if hasattr(obj, "read"):
            return self.read_buffer(obj, attrs=attrs)
        if isinstance(obj, pd.DataFrame):
            return self.read_dataframe(obj, attrs=attrs)
        if isinstance(obj, os.PathLike):
            if str(obj).endswith(".zip"):
                return self.read_zip(obj, attrs=attrs)
            elif ".tar" in str(obj):
                return self.read_tar(obj, attrs=attrs)
            return self.read_file_path(obj, attrs=attrs)
        if isinstance(obj, str):
            # See if this might be a file (make sure to expand user)
            if os.path.isfile(os.path.expanduser(obj)):
                p = Path(obj).expanduser()
                if p.suffix == ".zip":
                    return self.read_zip(p, attrs=attrs)
                elif p.suffix in (".tar", "tar.gz", "tar.bz"):
                    return self.read_tar(p, attrs=attrs)
                return self.read_file_path(p, attrs=attrs)
            if obj.startswith("http://") or obj.startswith("https://"):
                return self.read_url(obj, attrs=attrs)
            if obj.startswith("ftp://"):
                return self.read_ftp(obj, attrs=attrs)
            return self.read_string(obj, attrs=attrs)
        if isinstance(obj, bytes):
            return self.read_bytes(obj, attrs=attrs)
        raise ValueError(f"Could not read neuron from object of type '{type(obj)}'")

    def read_any_multi(
        self,
        objs,
        include_subdirs=DEFAULT_INCLUDE_SUBDIRS,
        parallel="auto",
        attrs: Optional[Dict[str, Any]] = None,
    ) -> "core.NeuronList":
        """Attempt to convert an arbitrary object into a NeuronList,
        potentially in parallel.

        Parameters
        ----------
        obj :               sequence
                            Sequence of anything readable by read_any_single or
                            directory path(s).
        include_subdirs :   bool
                            Whether to include subdirectories of a given
                            directory.
        parallel :          str | bool | int | None
                            "auto" or True for n_cores // 2, otherwise int for
                            number of jobs, or False for serial.
        attrs :             dict or None
                            Arbitrary attributes to include in each TreeNeuron
                            of the NeuronList.

        Returns
        -------
        core.NeuronList

        """
        if not utils.is_iterable(objs):
            objs = [objs]

        if not objs:
            logger.warning("No files found, returning empty NeuronList")
            return self.format_output(objs)

        new_objs = []
        for obj in objs:
            try:
                if is_dir(obj):
                    new_objs.extend(self.files_in_dir(obj, include_subdirs))
                    continue
            except TypeError:
                pass
            new_objs.append(obj)

        # `parallel` can be ("auto", threshold) in which case `threshold`
        # determines at what length we use parallel processing
        if isinstance(parallel, tuple):
            parallel, threshold = parallel
        else:
            threshold = 200

        if (
            isinstance(parallel, str)
            and parallel.lower() == "auto"
            and len(new_objs) < threshold
        ):
            parallel = False

        read_fn = partial(self.read_any_single, attrs=attrs)
        neurons = parallel_read(read_fn, new_objs, parallel)
        return self.format_output(neurons)

    def read_any(
        self,
        obj,
        include_subdirs=DEFAULT_INCLUDE_SUBDIRS,
        parallel="auto",
        limit=None,
        attrs: Optional[Dict[str, Any]] = None,
    ) -> "core.NeuronObject":
        """Attempt to read an arbitrary object into a neuron.

        Parameters
        ----------
        obj :       typing.IO | str | os.PathLike | pandas.DataFrame
                    Buffer, path to file or directory, URL, string, or
                    dataframe.

        Returns
        -------
        core.NeuronObject
        """
        if utils.is_iterable(obj) and not hasattr(obj, "read"):
            return self.read_any_multi(obj, parallel, include_subdirs, attrs=attrs)
        else:
            try:
                if is_dir(obj):
                    return self.read_directory(
                        obj, include_subdirs, parallel, limit, attrs=attrs
                    )
            except TypeError:
                pass
            try:
                if os.path.isfile(os.path.expanduser(obj)) and str(obj).endswith(
                    ".zip"
                ):
                    return self.read_zip(obj, parallel, limit, attrs=attrs)
                if os.path.isfile(os.path.expanduser(obj)) and ".tar" in str(obj):
                    return self.read_tar(obj, parallel, limit, attrs=attrs)
                if isinstance(obj, str) and obj.startswith("ftp://"):
                    return self.read_ftp(obj, parallel, limit, attrs=attrs)
            except TypeError:
                pass
            return self.read_any_single(obj, attrs=attrs)

    def parse_filename(self, filename: str) -> dict:
        """Extract properties from filename according to specified formatter.

        Parameters
        ----------
        filename : str

        Returns
        -------
        props :     dict
                    Properties extracted from filename.

        """
        # Make sure we are working with the filename not the whole path
        filename = Path(filename).name

        # Escape all special characters
        fmt = re.escape(self.fmt)

        # Unescape { and }
        fmt = fmt.replace("\\{", "{").replace("\\}", "}")

        # Replace all e.g. {name} with {.*}
        prop_names = []
        for prop in re.findall("{.*?}", fmt):
            prop_names.append(prop[1:-1].replace(" ", ""))
            fmt = fmt.replace(prop, "(.*)")

        # Match
        match = re.search(fmt, filename)

        if not match:
            raise ValueError(f'Unable to match "{self.fmt}" to filename "{filename}"')

        props = {"file": filename}
        for i, prop in enumerate(prop_names):
            for p in prop.split(","):
                # Ignore empty ("{}")
                if not p:
                    continue

                # If datatype was specified
                if ":" in p:
                    p, dt = p.split(":")

                    props[p] = match.group(i + 1)

                    if dt == "int":
                        props[p] = int(props[p])
                    elif dt == "float":
                        props[p] = float(props[p])
                    elif dt == "bool":
                        props[p] = bool(props[p])
                    elif dt == "str":
                        props[p] = str(props[p])
                    else:
                        raise ValueError(
                            f'Unable to interpret datatype "{dt}" ' f"for property {p}"
                        )
                else:
                    props[p] = match.group(i + 1)
        return props

    def _extract_connectors(self, nodes: pd.DataFrame) -> Optional[pd.DataFrame]:
        """Infer outgoing/incoming connectors from data.

        Parameters
        ----------
        nodes :     pd.DataFrame

        Returns
        -------
        Optional[pd.DataFrame]
                    With columns `["node_id", "x", "y", "z", "connector_id", "type"]`
        """
        return