nmx_io

class NMLReader(base.BaseReader):
    def __init__(
        self,
        precision: int = DEFAULT_PRECISION,
        attrs: Optional[Dict[str, Any]] = None,
        errors: str = "raise",
    ):
        super().__init__(
            fmt="",
            attrs=attrs,
            file_ext=".nml",
            read_binary=False,
            errors=errors,
            name_fallback="NML",
        )

        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 .nml buffer into a TreeNeuron.

        NML files are XML-encoded files containing data for a single neuron.

        Parameters
        ----------
        f :         IO
                    Readable buffer (must be bytes).
        attrs :     dict | None
                    Arbitrary attributes to include in the TreeNeuron.

        Returns
        -------
        core.TreeNeuron
        """
        return self.read_nml(f.read(), attrs=attrs)

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

        NML files are XML files containing a single neuron.

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

        Returns
        -------
        core.TreeNeuron
        """
        if isinstance(f, bytes):
            f = f.decode()

        f = io.StringIO(f)
        root = ET.parse(f).getroot()

        # Copy the attributes dict
        for element in root:
            if element.tag == "thing":
                nodes = pd.DataFrame.from_records([n.attrib for n in element[0]])
                edges = pd.DataFrame.from_records([n.attrib for n in element[1]])
                edges = edges.astype(self._dtypes["node_id"])

                nodes.rename({"id": "node_id"}, axis=1, inplace=True)
                nodes = nodes.astype(
                    {k: v for k, v in self._dtypes.items() if k in nodes.columns}
                )

        G = nx.Graph()
        G.add_edges_from(edges.values)
        tree = nx.bfs_tree(G, list(G.nodes)[0])
        edges = pd.DataFrame(list(tree.edges), columns=["source", "target"])
        nodes["parent_id"] = (
            edges.set_index("target").reindex(nodes.node_id.values).source.values
        )
        nodes["parent_id"] = nodes.parent_id.fillna(-1).astype(self._dtypes["node_id"])
        nodes.sort_values("node_id", inplace=True)

        return core.TreeNeuron(
            nodes, **(self._make_attributes({"name": "NML", "origin": "nml"}, attrs))
        )

class NMXReader(NMLReader):
    """This is a version of the NML file reader that reads from zipped archives."""

    def __init__(
        self,
        precision: int = DEFAULT_PRECISION,
        attrs: Optional[Dict[str, Any]] = None,
        errors: str = "raise",
    ):
        super().__init__(precision=precision, errors=errors, attrs=attrs)

        # Overwrite some of the settings
        self.read_binary = True
        self.file_ext = ".nmx"
        self.name_fallback = "NMX"

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

        NMX files are zip files containing XML-encoded .nml files containing
        data for a single neuron.

        Parameters
        ----------
        f :         IO
                    Readable buffer (must be bytes).
        attrs :     dict | None
                    Arbitrary attributes to include in the TreeNeuron.

        Returns
        -------
        core.TreeNeuron
        """
        if not isinstance(f.read(0), bytes):
            raise ValueError(f'Expected bytes, got "{type(f.read(0))}"')

        zip = ZipFile(f)
        for f in zip.filelist:
            if f.filename.endswith(".nml") and "skeleton" in f.filename:
                attrs["file"] = f.filename
                attrs["id"] = f.filename.split("/")[0]
                return self.read_nml(zip.read(f), attrs=attrs)
        logger.warning(
            f'Skipped "{f.filename.split("/")[0]}.nmx": failed to ' 'import skeleton.'
        )