update_hi.py

#!/usr/bin/env python3
# -*- coding: utf-8; mode: python -*-

"""
After making changes to the JSON file, update the CSV and write the
encoded string into the TEI file.
"""

__version__ = "1.0"
__date__ = "20191127"
__author__ = "kthoden@mpiwg-berlin.mpg.de"

import argparse
import logging
import base64
import sys
import json
import csv
import os
import shutil
from pathlib import Path
from lxml import etree

ns_tei = "http://www.tei-c.org/ns/1.0"
NS_MAP = {"t" : ns_tei}

logging.basicConfig(level=logging.DEBUG, format=' %(asctime)s - %(levelname)s - %(message)s')

DEFAULT_OUTPUT_DIR = \
    Path(os.environ['OUTPUT_DIR'] if 'OUTPUT_DIR' in os.environ else './output')

def load_json_file(fig_id, input_dir):
    """Load a json file. Return it formatted."""

    filename = f"{input_dir}/{fig_id}.json"

    logging.debug("Opening %s.", filename)

    with open(filename, "r") as jsonfile:
        json_string = json.load(jsonfile)

    js_json = json.dumps(json_string)

    return js_json
# def load_json_file ends here


def base64_encode(json):
    """Encode the JSON string"""

    byte_string = json.encode()
    encoded_json = base64.encodebytes(byte_string)
    decoded_json = encoded_json.decode('utf-8')
    encoded_json_single_line = decoded_json.replace("\n", "")

    return encoded_json_single_line
# def base64_encode ends here


def update_xml_entry(xml_file, fig_id, json_encoded):
    """Update the entry in the XML file"""

    backup_file = str(xml_file).replace(".xml", "-backup.xml")
    shutil.copy(xml_file, backup_file)

    xml_tree = etree.parse(str(xml_file))
    entry_to_replace = xml_tree.xpath(f"//t:div[@xml:id='{fig_id}']/t:ab", namespaces=NS_MAP)[0]

    entry_to_replace.text = ""
    entry_to_replace.text = json_encoded

    xml_tree.write(str(xml_file), pretty_print=True, xml_declaration=True,encoding="utf-8")
    logging.info(f"Wrote XML file: {xml_file}.")
# def update_xml_entry ends here

def update_csv_entry(csv_file, fig_id, json_encoded):
    """Update the line in the CSV file"""

    # Thanks to https://stackoverflow.com/a/46130947
    import tempfile

    tmp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
    fieldnames = ['xmlid', 'hiid', 'layers', 'elementstring']

    backup_file = str(csv_file).replace(".csv", "-backup.csv")
    shutil.copy(csv_file, backup_file)

    with open(csv_file, 'r') as csvfile, tmp_file:
        reader = csv.DictReader(csvfile, fieldnames=fieldnames)
        writer = csv.DictWriter(tmp_file, fieldnames=fieldnames)
        for row in reader:
            if row['xmlid'] == str(fig_id):
                logging.info('Updating entry of %s', fig_id)
                row = {'xmlid': row['xmlid'], 'hiid': row['hiid'], 'layers': row['layers'], 'elementstring': json_encoded}
                writer.writerow(row)
            else:
                writer.writerow(row)

    shutil.move(tmp_file.name, csv_file)
# def update_csv_entry ends here


def main():
    """The main bit"""

    parser = argparse.ArgumentParser()
    parser.add_argument("figid", help="The xml:id of the figure whose annotation is changed.")
    parser.add_argument("-c", "--csv", help="Use CSV input instead of XML", action="store_true")
    parser.add_argument(
            "-o", "--output-dir",
            default = DEFAULT_OUTPUT_DIR / "../hyperimage",
            help="output directory"
    )
    args = parser.parse_args()

    OUTPUT_DIR = Path( args.output_dir )
    JSON_OUTPUT_DIR = OUTPUT_DIR / "json"
    json_string = load_json_file(args.figid, JSON_OUTPUT_DIR)
    encoded_json = base64_encode(str(json_string))

    if args.csv:
        CSV_FILE = OUTPUT_DIR / "hi_figures.csv"
        update_csv_entry(CSV_FILE, args.figid, encoded_json)
    else:
        XML_FILE = OUTPUT_DIR / "hi_figures.xml"
        update_xml_entry(XML_FILE, args.figid, encoded_json)
    logging.info("Successfully updated %s", args.figid)
# def main ends here

if __name__ == '__main__':
    main()
# finis