Merge pull request #185 from NeuralEnsemble/feat/more-docs

More documentation for writers/loaders

neuroml/loaders.py

@@ -1,23 +1,30 @@
-from neuroml.nml.nml import parse as nmlparse
-
-from neuroml.nml.nml import parseString as nmlparsestring
-
-import neuroml
-from neuroml import (NeuroMLDocument)
-import neuroml.utils as utils
-
 import os
 import sys
 import warnings
-
 from typing import Callable, Optional
 
+import neuroml
+import neuroml.utils as utils
+from neuroml import NeuroMLDocument
+from neuroml.nml.nml import parse as nmlparse
+from neuroml.nml.nml import parseString as nmlparsestring
+
 supressGeneratedsWarnings = True
 
 
 class NeuroMLLoader(object):
+    """Class for loading NeuroML."""
+
     @classmethod
-    def load(cls, src):
+    def load(cls, src: str) -> neuroml.NeuroMLDocument:
+        """Load a NeuroML file.
+
+        :param src: file
+        :type src: str
+        :returns: NeuroMLDocument object
+        :rtype: neuroml.NeuromlDocument
+        :raises TypeError: if the file is not a valid NeuroML document
+        """
         doc = cls.__nml2_doc(src)
         if isinstance(doc, NeuroMLDocument):
             return doc
@@ -29,7 +36,15 @@ def load(cls, src):
             )
 
     @classmethod
-    def __nml2_doc(cls, file_name):
+    def __nml2_doc(cls, file_name: str) -> neuroml.NeuroMLDocument:
+        """Load and parse a NeuroML file.
+
+        :param file_name: the file
+        :type file_name: str
+        :returns: NeuroMLDocument object
+        :rtype: neuroml.NeuromlDocument
+        :raises Exception: if the document is not a valid NeuroML Document
+        """
         try:
             if supressGeneratedsWarnings:
                 warnings.simplefilter("ignore")
@@ -43,13 +58,43 @@ def __nml2_doc(cls, file_name):
 
 
 class NeuroMLHdf5Loader(object):
+    """Class for loading a NeuroML HDF5 file."""
+
     @classmethod
-    def load(cls, src, optimized=False):
+    def load(cls, src: str, optimized: bool = False) -> neuroml.NeuroMLDocument:
+        """Load a NeuroML HDF5 file.
+
+        :param src: file
+        :type src: str
+        :param optimized: load optimized numpy representation
+            In the optimized representation, instead of the complete Python
+            object tree being constructed, the various tables in the HDF5 file
+            are loaded as numpy arrays. This is transparent to the user, who
+            can continue using the standard methods to access the data.
+        :type optimized: bool
+        :returns: NeuroMLDocument object
+        :rtype: neuroml.NeuromlDocument
+        """
         doc = cls.__nml2_doc(src, optimized)
         return doc
 
     @classmethod
-    def __nml2_doc(cls, file_name, optimized=False):
+    def __nml2_doc(
+        cls, file_name: str, optimized: bool = False
+    ) -> neuroml.NeuroMLDocument:
+        """Load and parse a NeuroML HDF5 file.
+
+        :param file_name: the file
+        :type file_name: str
+        :param optimized: load optimized numpy representation
+            In the optimized representation, instead of the complete Python
+            object tree being constructed, the various tables in the HDF5 file
+            are loaded as numpy arrays. This is transparent to the user, who
+            can continue using the standard methods to access the data.
+        :type optimized: bool
+        :returns: NeuroMLDocument object
+        :rtype: neuroml.NeuromlDocument
+        """
         import logging
 
         logging.basicConfig(
@@ -96,6 +141,7 @@ def load_swc_single(cls, src, name=None):
         )
 
         import numpy as np
+
         from neuroml import arraymorph
 
         dtype = {

neuroml/writers.py

@@ -1,6 +1,7 @@
-import neuroml
-from six import string_types
+import typing
 
+import neuroml
+from neuroml.arraymorph import ArrayMorphology
 
 """Classes to write NeuroML to various formats."""
 
@@ -10,8 +11,9 @@ class NeuroMLWriter(object):
 
     In future can implement from other types via chain of responsibility pattern.
     """
+
     @classmethod
-    def write(cls, nmldoc, file, close=True):
+    def write(cls, nmldoc: neuroml.NeuroMLDocument, file: str, close: bool = True):
         """Write a NeuroMLDocument to file.
 
         :param nmldoc: NeuroML document object to write
@@ -23,8 +25,10 @@ def write(cls, nmldoc, file, close=True):
         :raises AttributeError: if export fails
         """
 
-        if isinstance(file, string_types):
-            file = open(file, "w")
+        if isinstance(file, str):
+            fileh = open(file, "w")
+        else:
+            fileh = file
 
         # TODO: this should be extracted from the schema:
         namespacedef = 'xmlns="http://www.neuroml.org/schema/neuroml2" '
@@ -37,20 +41,27 @@ def write(cls, nmldoc, file, close=True):
 
         try:
             nmldoc.export(
-                file, 0, name_="neuroml", namespacedef_=namespacedef
+                fileh, 0, name_="neuroml", namespacedef_=namespacedef
             )  # name_ param to ensure root element named correctly - generateDS limitation
         except AttributeError as ae:
-            file.close()
+            fileh.close()
             raise (ae)
 
         if close:
-            file.close()
+            fileh.close()
 
 
 class NeuroMLHdf5Writer(object):
     """Exports NeuroML documents to HDF5 format."""
+
     @classmethod
-    def write(cls, nml_doc, h5_file_name, embed_xml=True, compress=True):
+    def write(
+        cls,
+        nml_doc: neuroml.NeuroMLDocument,
+        h5_file_name: str,
+        embed_xml: bool = True,
+        compress: bool = True,
+    ):
         """Write a NeuroMLDocument to HDF5 file
 
         :param nmldoc: NeuroML document object to write
@@ -92,9 +103,11 @@ def write(cls, nml_doc, h5_file_name, embed_xml=True, compress=True):
 
             try:
                 import StringIO
+
                 sf = StringIO.StringIO()
             except ImportError:
                 import io
+
                 sf = io.StringIO()
 
             NeuroMLWriter.write(nml_doc, sf, close=False)
@@ -140,13 +153,18 @@ class ArrayMorphWriter(object):
     """
 
     @classmethod
-    def __write_single_cell(cls, array_morph, fileh, cell_id=None):
+    def __write_single_cell(
+        cls,
+        array_morph: ArrayMorphology,
+        fileh,
+        cell_id: typing.Optional[str] = None,
+    ):
         """Write a array morphology to a file handler.
 
         :param array_morph: a array morph object containing a morphology
-        :type array_morph: neuroml.arraymorph.ArrayMorphology
-        :param fileh: file handler of file to write to
-        :type fileh: file object
+        :type array_morph: ArrayMorphology
+        :param fileh: pytables file object of file to write to
+        :type fileh: pytables file object
         :param cell_id: id of cell
         :type cell_id: str
         """
@@ -182,7 +200,7 @@ def __write_single_cell(cls, array_morph, fileh, cell_id=None):
         )
 
     @classmethod
-    def __write_neuroml_document(cls, document, fileh):
+    def __write_neuroml_document(cls, document: neuroml.NeuroMLDocument, fileh):
         """Write a NeuroMLDocument containing morphology to a file handler
 
         :param document: a NeuroML document object containing a morphology
@@ -207,11 +225,15 @@ def __write_neuroml_document(cls, document, fileh):
             cls.__write_single_cell(morphology, fileh, cell_id=cell.id)
 
     @classmethod
-    def write(cls, data, filepath):
+    def write(
+        cls,
+        data: typing.Union[neuroml.NeuroMLDocument, ArrayMorphology],
+        filepath: str,
+    ):
         """Write morphology to file in ArrayMorph format.
 
         :param data: data to write
-        :type data: neuroml.arraymorph.ArrayMorphology or neuroml.NeuroMLDocument
+        :type data: ArrayMorphology or neuroml.NeuroMLDocument
         :param filepath: path of file to write to
         :type filepath: str
 
@@ -223,7 +245,7 @@ def write(cls, data, filepath):
         # Now instead we should go through a document/cell/morphology
         # hierarchy - this kind of tree traversal should be done recursively
 
-        if isinstance(data, neuroml.arraymorph.ArrayMorphology):
+        if isinstance(data, ArrayMorphology):
             cls.__write_single_cell(data, fileh)
 
         if isinstance(data, neuroml.NeuroMLDocument):