swc_io

class SwcReader(base.BaseReader):
    def __init__(
        self,
        connector_labels: Optional[Dict[str, Union[str, int]]] = None,
        soma_label: Union[str, int] = 1,
        delimiter: str = DEFAULT_DELIMITER,
        precision: int = DEFAULT_PRECISION,
        read_meta: bool = False,
        fmt: str = DEFAULT_FMT,
        errors: str = "raise",
        attrs: Optional[Dict[str, Any]] = None,
    ):
        if not fmt.endswith(".swc"):
            raise ValueError('`fmt` must end with ".swc"')

        super().__init__(
            fmt=fmt, attrs=attrs, file_ext=".swc", errors=errors, name_fallback="SWC"
        )
        self.connector_labels = connector_labels or dict()
        self.soma_label = soma_label
        self.delimiter = delimiter
        self.read_meta = read_meta

        int_, float_ = base.parse_precision(precision)
        self._dtypes = {
            "node_id": int_,
            "parent_id": int_,
            "label": "category",
            "x": float_,
            "y": float_,
            "z": float_,
            "radius": float_,
        }

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

        Parameters
        ----------
        f :         IO
                    Readable buffer (if bytes, interpreted as utf-8).
        attrs :     dict | None
                    Arbitrary attributes to include in the TreeNeuron.

        Returns
        -------
        core.TreeNeuron
        """
        if isinstance(f, HTTPResponse):
            f = io.StringIO(f.data.decode())

        if isinstance(f.read(0), bytes):
            f = io.TextIOWrapper(f, encoding="utf-8")

        header_rows = read_header_rows(f)
        try:
            nodes = pd.read_csv(
                f,
                delimiter=self.delimiter,
                skipinitialspace=True,
                skiprows=len(header_rows),
                comment=COMMENT,
                header=None,
                na_values=NA_VALUES,
            )
            if len(nodes.columns) < len(NODE_COLUMNS):
                raise ValueError("Not enough columns in SWC file.")
            elif len(nodes.columns) > len(NODE_COLUMNS):
                logger.warning(
                    f"Found {len(nodes.columns)} instead of the expected 7 "
                    "columns in SWC file. Assuming additional columns are "
                    "custom properties. You can silence this warning by setting "
                    "`navis.set_loggers('ERROR')`."
                )
                nodes.columns = (
                    list(NODE_COLUMNS) + nodes.columns[len(NODE_COLUMNS) :].tolist()
                )
            else:
                nodes.columns = NODE_COLUMNS
        except pd.errors.EmptyDataError:
            # If file is totally empty, return an empty neuron
            # Note that the TreeNeuron will still complain but it's a better
            # error message
            nodes = pd.DataFrame(columns=NODE_COLUMNS)

        # Check for row with JSON-formatted meta data
        # Expected format '# Meta: {"id": "12345"}'
        if self.read_meta:
            meta_row = [r for r in header_rows if r.lower().startswith("# meta:")]
            if meta_row:
                meta_data = json.loads(meta_row[0][7:].strip())
                attrs = base.merge_dicts(meta_data, attrs)

        return self.read_dataframe(
            nodes, base.merge_dicts({"swc_header": "\n".join(header_rows)}, attrs)
        )

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

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

        Returns
        -------
        core.TreeNeuron
        """
        n = core.TreeNeuron(
            sanitise_nodes(nodes.astype(self._dtypes, errors="ignore", copy=False)),
            connectors=self._extract_connectors(nodes),
        )

        if self.soma_label is not None:
            is_soma_node = n.nodes.label.values == self.soma_label
            if any(is_soma_node):
                n.soma = n.nodes.node_id.values[is_soma_node][0]

        attrs = self._make_attributes({"name": "SWC", "origin": "DataFrame"}, attrs)

        # SWC is special - we do not want to register it
        n.swc_header = attrs.pop("swc_header", "")

        # Try adding properties one-by-one. If one fails, we'll keep track of it
        # in the `.meta` attribute
        meta = {}
        for k, v in attrs.items():
            try:
                n._register_attr(k, v)
            except (AttributeError, ValueError, TypeError):
                meta[k] = v

        if meta:
            n.meta = meta

        return n

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

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

        Returns
        -------
        Optional[pd.DataFrame]
                    With columns `["node_id", "x", "y", "z", "connector_id", "type"]`
        """
        if not self.connector_labels:
            return None

        to_concat = [
            pd.DataFrame([], columns=["node_id", "connector_id", "type", "x", "y", "z"])
        ]
        for name, val in self.connector_labels.items():
            cn = nodes[nodes.label == val][["node_id", "x", "y", "z"]].copy()
            cn["connector_id"] = None
            cn["type"] = name
            to_concat.append(cn)

        return pd.concat(to_concat, axis=0)