Source code for SDF.file_io.sdf

from textwrap import wrap
from typing import Union, Optional, TextIO
from xml.etree.ElementTree import XMLParser, parse, Element
from xml.sax.saxutils import escape, quoteattr

from SDF.data_model import Dataset, SDFObject, Workspace
from SDF.sdf_rc import ENCODING, TABSIZE, LINEWIDTH


def load_from_sdf(filename: Union[str, TextIO]) -> SDFObject:
    parser = XMLParser(encoding=ENCODING)
    root = parse(filename, parser=parser).getroot()
    if root.tag == "workspace":
        return Workspace.from_xml_element(root)
    elif root.tag == "dataset":
        return Dataset.from_xml_element(root)
    else:
        raise ValueError(f"Expected the root to be a <dataset> or <workspace>, got {root.tag}")


def write_sdf(obj: SDFObject, file: Union[str, TextIO]) -> None:
    if isinstance(obj, Workspace):
        doctype = "workspace"
    elif isinstance(obj, Dataset):
        doctype = "dataset"
    else:
        raise TypeError(f"Can only write datasets and workspaces, got {obj}")

    if isinstance(file, str):
        fp = open(file, "w")
    elif hasattr(file, "write"):
        fp = file
    else:
        raise ValueError("Given file is not a filename and not writable")

    fp.write(f'<?xml version="1.0" encoding="{ENCODING}"?>\n')
    fp.write(f'<!DOCTYPE {doctype}>\n')
    __xml_to_file(obj.to_xml_element(), fp)

    if isinstance(file, str):
        fp.close()


def __xml_to_file(element: Optional[Element], file: TextIO, indent_level: int = 0) -> None:
    """Writes an XML string of this object to a given file"""
    if element is None:
        return

    indent = " " * indent_level * TABSIZE

    # opening tag
    file.write(indent)
    file.write(f"<{element.tag}")
    for key, value in element.attrib.items():
        file.write(f" {escape(key.strip())}={quoteattr(value.strip())}")
    if len(element) == 0 and (element.text is None or element.text.isspace()):
        file.write("/>\n")
        return
    else:
        file.write(">\n")

    # text
    __text_to_file(element.text, file, indent_level+1)

    # children
    for child in element:
        __xml_to_file(child, file, indent_level+1)

    # closing tag
    file.write(indent)
    file.write(f"</{element.tag}>\n")


def __text_to_file(text: Optional[str], file: TextIO, indent_level: int = 0):
    """Writes a given string to a given file. Respects indentation and wraps long lines"""
    if text is None:
        return

    indent = " " * indent_level * TABSIZE

    for line in text.splitlines(keepends=False):
        if len(indent) + len(line) <= LINEWIDTH:
            file.write(indent + line + "\n")
        else:
            for wline in wrap(line, width=LINEWIDTH, initial_indent=indent, break_on_hyphens=False,
                              subsequent_indent=indent, break_long_words=False, tabsize=TABSIZE):
                file.write(wline + "\n")