cdm_data_loaders.readers.xml ¶

def parse_head_matter(file_path: FileSource | IO[bytes]) -> dict[str, str]:
    """Parse the namespace declarations from the head matter of an XML file.

    Reads the beginning of an XML document (gzipped or plain) and collects all XML
    namespace declarations (``xmlns`` / ``xmlns:prefix`` attributes) that are in scope
    on the root element. Parsing stops as soon as the root element is fully opened, so
    the whole document is never loaded into memory.

    The returned mapping uses namespace prefixes as keys and namespace URIs as values.
    The default (unprefixed) namespace is stored under the empty-string key ``""``.

    Given an open handle, parsing starts at the handle's current position and the handle is
    left open, positioned wherever parsing stopped; call ``seek(0)`` before reading it again.

    :param file_path: local path, dlt file item, or open binary handle; the file can be gzipped or not.
    :type file_path: FileSource | IO[bytes]
    :raises FileNotFoundError: if the file does not exist.
    :raises lxml.etree.XMLSyntaxError: if the file is empty or malformed.
    :return: mapping of namespace prefix to namespace URI.
    :rtype: dict[str, str]
    """
    logger.debug("Parsing head matter from %s", source_url(file_path))

    namespaces: dict[str, str] = {}
    with open_binary(file_path) as fh:
        ctx = iterparse(fh, events=("start-ns", "start"))
        for event, elem in ctx:
            if event == "start-ns":
                # elem is a (prefix, uri) tuple; the default namespace has an empty prefix.
                prefix, uri = elem
                namespaces[prefix] = uri
                continue

            # The first "start" event corresponds to the root element; by this point every
            # namespace declared on the root has already been emitted as a "start-ns" event.
            break

    logger.debug("Found %d namespace declaration(s) in %s", len(namespaces), source_url(file_path))
    return namespaces

def stream_xml_file(file_path: FileSource | IO[bytes], element_with_ns: str) -> Generator[Element, Any]:
    """Stream the outermost matching XML elements from a file.

    A matching element nested inside another matching element is yielded as part of the outer element.

    Given an open handle, streaming starts at the handle's current position and the handle is left open.

    :param file_path: local path, dlt file item, or open binary handle; the file can be gzipped or not.
    :type  file_path: FileSource | IO[bytes]
    :param element_with_ns: name of the element (including namespace, in braces) to return; e.g. f"{{{UNIPROT_NS}}}entry"
    :type  element_with_ns: str
    :yield: elements from the file
    :rtype: Generator[Element, Any]
    """
    logger.debug("Streaming XML from %s", source_url(file_path))

    with open_binary(file_path) as f:
        # number of open matching elements; 0 means the parser is outside any match
        depth = 0
        for event, elem in iterparse(f, events=("start", "end"), tag=element_with_ns, remove_blank_text=True):
            if event == "start":
                depth += 1
                continue
            depth -= 1
            if depth:
                # nested match: leave it in place as part of the enclosing match
                continue
            yield elem
            _clear_element(elem)

def stream_xml_file_at_level(file_path: FileSource | IO[bytes], level: int) -> Generator[Element, Any]:
    """Stream all XML elements at depth ``level`` relative to the root element.

    Depth ``0`` yields the root element; depth ``1`` yields all direct children of the root
    element; depth ``2`` yields all grandchildren of the root element, and so on.

    Given an open handle, streaming starts at the handle's current position and the handle is left open.

    :param file_path: local path, dlt file item, or open binary handle; the file can be gzipped or not.
    :type  file_path: FileSource | IO[bytes]
    :param level: non-negative depth relative to the root element
    :type  level: int
    :yield: elements at depth ``level`` from the file
    :rtype: Generator[Element, Any]
    """
    if level < 0:
        msg = f"level must be >= 0, got {level}"
        raise ValueError(msg)

    logger.debug("Streaming XML at level %d from %s", level, source_url(file_path))

    with open_binary(file_path) as f:
        current_depth = -1
        for event, elem in iterparse(f, events=("start", "end"), remove_blank_text=True):
            if event == "start":
                current_depth += 1
                continue
            if current_depth == level:
                yield elem
                _clear_element(elem)
            elif 0 < current_depth < level:
                _clear_element(elem)
            current_depth -= 1

def process_xml_file_to_pyarrow(file_path: FileSource, settings: XmlToPyArrowSettings) -> Generator[DataItemWithMeta]:
    """Convert XML elements to PyArrow table batches using an XSD schema.

    When ``settings.xml_level`` is set, all elements at that depth are streamed, converted using
    ``settings.xsd_schema``, and routed to tables named after each element's local tag name
    (with PyArrow schemas from ``settings.schema_map``). Otherwise, elements matching
    ``settings.xml_tag`` are streamed and buffered into ``settings.table_name``.

    :param file_path: local path or dlt file item
    :type  file_path: FileSource
    :param settings: pipeline config, including xml_tag/table_name or xml_level, xsd_schema, schema_map, and buffer_size
    :type  settings: XmlToPyArrowSettings
    :yield: pages of PyArrow tables
    :rtype: Generator[DataItemWithMeta]
    """
    xml_level = getattr(settings, "xml_level", None)
    if isinstance(xml_level, int):
        yield from _process_xml_file_to_pyarrow_at_level(file_path, settings, xml_level)
        return

    if settings.table_name is None or settings.xml_tag is None:
        err_msg = "Expected table_name and xml_tag when xml_level is not set"
        raise RuntimeError(err_msg)

    convert_element = settings.xml_element_converter
    if convert_element is None:
        err_msg = "Expected xml_element_converter when xml_level is not set"
        raise RuntimeError(err_msg)

    yield from _process_single_tag_xml_file(
        file_path,
        settings,
        convert_element,
        resolve_default_ns=True,
    )

def process_xml_file_to_dict(file_path: FileSource, settings: XmlToDictSettings) -> Generator[DataItemWithMeta]:
    """Convert XML elements to dictionary form using ``xmltodict.parse``.

    :param file_path: local path or dlt file item
    :type  file_path: FileSource
    :param settings: pipeline config, including xml_tag, table_name, and buffer_size
    :type  settings: XmlToDictSettings
    :yield: pages of dictionary items
    :rtype: Generator[DataItemWithMeta]
    """
    yield from _process_single_tag_xml_file(
        file_path,
        settings,
        _parse_element_to_dict,
    )

def process_xml_file(
    file_path: FileSource,
    settings: CtsSettings,
    xml_tag: str,
    parse_fn: Callable,
) -> Generator[DataItemWithMeta, Any]:
    """Core generator shared by XML-based dlt pipeline resources.

    This processor is expected to return a dictionary of table names and lists of rows, unlike the
    xmltodict parser, which creates a single dictionary for each element.

    ``parse_fn`` receives the file's source id (see ``cdm_data_loaders.utils.file_items.source_id``) as
    ``file_path``: the item's relative path for a file item, or the file name for a local path. The
    recorded source file therefore does not depend on where the input lives.

    :param file_path: local path or dlt file item for the XML file to be processed
    :type  file_path: FileSource
    :param settings: pipeline config with buffer_size and log_interval
    :type  settings: CtsSettings
    :param xml_tag: XML element tag to stream
    :type  xml_tag: str
    :param parse_fn: callable(entry, file_path) -> dict[str, rows]
    :type  parse_fn: Callable
    :yield: table-tagged rows
    :rtype: Generator[DataItemWithMeta, Any]
    """
    buffer = make_dict_buffer(settings, DictBuffer)
    file_source_id = source_id(file_path)
    yield from _stream_and_buffer_xml_elements(
        file_path=file_path,
        xml_tag=xml_tag,
        log_interval=settings.log_interval,
        buffer=buffer,
        buffer_element_fn=lambda element: buffer.add_items(parse_fn(entry=element, file_path=file_source_id)),
    )

def build_xml_file_resource(
    settings: XmlIngestSettings,
    xml_tag: str,
    parse_fn: Callable,
    resource_name: str,
) -> DltResource:
    """Build a dlt resource that parses the XML files in `settings.input_dir`.

    Files are read using the dlt filesystem source.

    :param settings: pipeline config with input_dir, file_glob, buffer_size and log_interval
    :type  settings: XmlIngestSettings
    :param xml_tag: XML element tag to stream
    :type  xml_tag: str
    :param parse_fn: function for parsing the XML
    :type  parse_fn: Callable
    :param resource_name: name of the resulting dlt resource
    :type  resource_name: str
    :return: resource yielding table-tagged rows
    :rtype: DltResource
    """
    return build_per_file_filesystem_resource(
        settings,
        process_xml_file,
        xml_tag,
        parse_fn,
        name=resource_name,
        parallelized=False,
    )