Periodic-Table-JSON/scripts/specific_data.py

#!/usr/bin/env python3

"""Select specific data about elements and write them to a named file.

Use the --help for options and examples.

"""

import sys
import os
import argparse
import textwrap
import json
from pathlib import Path
from collections import namedtuple

ColorCodes = namedtuple('ColorCodes', ['PURPLE', 'CYAN', 'DARKCYAN',
                                       'BLUE', 'GREEN', 'YELLOW',
                                       'RED', 'BOLD', 'UNDERLINE',
                                       'END',])

color_code = ColorCodes(PURPLE='\033[95m', CYAN='\033[96m', DARKCYAN='\033[36m',
                         BLUE='\033[94m', GREEN='\033[92m', YELLOW='\033[93m',
                         RED='\033[91m', BOLD='\033[1m', UNDERLINE='\033[4m',
                         END='\033[0m',)

# pylint: disable=missing-function-docstring
def create_commandeline_parser(default_file):
    parser = argparse.ArgumentParser(
        description=__doc__,
        epilog=textwrap.dedent(f"""
        examples:
          Properties written to a json file:
             $ {sys.argv[0]} --properties=name,atomic_mass --output name_mass.json

          Properties written to a csv file:
             $ {sys.argv[0]} --properties name,atomic_mass --output name_mass.csv

          Properties written into both files {default_file}.json and {default_file}.csv:
             $ {sys.argv[0]} --properties=name,atomic_mass

          Union of properties written into both files {default_file}.json and {default_file}.csv:
             $ {sys.argv[0]} --properties=name,atomic_mass --interactive

          Select properties interactively and write to files {default_file}.json and {default_file}.csv:
             $ {sys.argv[0]} --interactive

        NOTE: output files are written to the directory above {sys.argv[0]}.

        """),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )

    def comma_separated(arg):
        return arg.split(",")
    parser.add_argument('--properties', type=comma_separated, metavar='P1,...', nargs=1,
                        help='comma separated list of properties')

    parser.add_argument('--interactive', action="store_true",
                        help='whether to interactively select data')

    parser.add_argument('--output', metavar='FILENAME', nargs='?', const=default_file, default='',
                        help=f'where to output the data (default: {default_file}.{{json,csv}})')
    return parser


def read_periodic_table():
    """
    Load and return elements from the Periodic Table JSON file.

    This function reads the Periodic Table data from a JSON file located
    one directory above the current file's directory.

    Returns:
        tuple: A tuple containing two elements:
            - A list of dictionaries, where each dictionary represents an element
              with its properties as key-value pairs.
            - A set of strings, representing the keys (properties) available in
              the element dictionaries.
    """
    with open(os.path.join(Path(__file__).parents[1],'PeriodicTableJSON.json'),
              encoding="utf8") as file:
        elements = json.load(file)['elements']
    return elements, elements[0].keys()


def parse_properties(data_needed, args, keys):
    """
    Update data_needed with properties specified in args if they exist in keys.

    Parameters:
        data_needed (dict): A dictionary where each key represents a property name,
            and the value indicates whether the property is needed (True or False).
        args (argparse.Namespace): An object containing command line arguments.
            This function expects args to have a 'properties' attribute, which is
            a list of lists containing property names as strings.
        keys (list): A list of strings representing valid property names.

    Returns:
        dict: The updated data_needed dictionary with keys set to True for properties
           that are both requested in args and exist in keys.

    Side Effects:
        - Prints messages to the console indicating whether each requested property
          was found or not found.
        - Modifies the data_needed dictionary in-place by updating the values for
          keys that represent valid properties.
    """
    def show_bad(pval):
        message = ''.join([color_code.RED, 'Property ',
                           color_code.BLUE, pval,
                           color_code.RED, ' not found.',
                           color_code.END,])
        print(message)

    def show_good(pval):
        message = ''.join([color_code.GREEN, 'Property ',
                           color_code.BLUE, pval,
                           color_code.GREEN,' found.',
                           color_code.END,])
        print(message)

    props = set(args.properties[0])
    keys_as_set = set(keys)
    bad = props - keys_as_set
    good = props & keys_as_set
    data_needed.update({k:True for k in good})

    for pval in bad:
        show_bad(pval)
    for pval in good:
        show_good(pval)
    return data_needed


def parse_interactive(data_needed, keys):
    """
    Interactively update the data_needed dictionary based on user input.

    This function iterates over a list of keys, presenting each to the user
    and asking if it is needed. The user's response is used to update the
    data_needed dictionary: setting the value to True for keys that are needed.

    Parameters:
        data_needed (dict): A dictionary where each key represents a potential
            option, and the value (True/False) indicates whether it is needed.
        keys (list): A list of strings representing all possible options that
            the user can choose from.

    Returns:
        dict: The updated data_needed dictionary reflecting the user's choices.
            Keys corresponding to needed options are set to True.


    Note:
        The function calls an external 'clear()' function to clear the console
        before each iteration, improving readability of the interactive session.
    """
    def show_selected():
        message = ''.join([
            color_code.BOLD, color_code.GREEN,
            f'{len(data_needed)}', ' Option(s) Selected ',
            color_code.END,
            color_code.UNDERLINE, f'{list(data_needed.keys())}',
            color_code.END,
        ])
        print(message)

    def show_next():
        message = ''.join([
            'Do you need ',
            color_code.BOLD, color_code.BLUE, key, '? ',
            color_code.END,
        ])
        print(message)

    def default_input(prompt, default='y'):
        user_input = input(prompt)
        return user_input if  user_input else default

    def select_next():
        prompt = ''.join([
            color_code.GREEN, 'y/',
            color_code.RED, 'n/',
            color_code.BLUE, 'q(uit)',
            color_code.END, '[',
            color_code.PURPLE, 'default: ',
            color_code.GREEN, 'y',
            color_code.END, ']:',
        ])
        return default_input(prompt)

    def clear():
        os.system('cls' if os.name == 'nt' else 'clear')

    clear()
    for key in keys:
        if key in data_needed:
            continue
        done = False
        while True:
            show_selected()
            show_next()
            needed = select_next()
            if needed == 'y':
                data_needed[key] = True
                break
            if needed == 'n':
                break
            if needed == 'q':
                done = True
                break
            print('Invalid input')
            continue
        if done:
            break
        clear()
    return data_needed


def save2file(args, elements, data_needed, default_file):
    """
    Save elements to file(s) in JSON and/or CSV format based on args.

    This function writes the specified elements to a file or files, using
    JSON and/or CSV formats.

    Parameters:
        args: An object containing command line arguments, including 'output'
              which specifies the output file name and format.
        elements (list): A list of dictionaries, each representing an element
                         with data to be saved.
        data_needed (dict): A dictionary specifying which data from elements
                            should be included in the output file(s).
        default_file (str): The default file name to use if no output file is
                            specified or if the specified file is the default.
    """
    if not args.output or args.output == default_file:
        # Use default and write to both csv and json files.
        write_json(default_file, elements, data_needed)
        write_csv(default_file, elements, data_needed)
    else:
        # Write to provided file name.
        if ('json' in args.output.lower()) or ('csv' in args.output.lower()):
            output = args.output.replace('.json', '').replace('.csv', '')
        if 'json' in args.output.lower():
            write_json(output, elements, data_needed)
        if 'csv' in args.output.lower():
            write_csv(output, elements, data_needed)


def write_csv(output, elements, data_needed):
    """
    Write selected data from elements to a CSV file.

    Parameters:
        output (str): The base name of the output file.
        elements (list): A list of dictionaries with element data.
        data_needed (dict): A dict indicating which keys to include.
    """
    with open(os.path.join(Path(__file__).parents[1], output + '.csv'),
              'w', encoding="utf8") as file:
        elem_to_write = []
        elem_to_write.append(','.join(data_needed.keys()))
        for element in elements:
            elmnt = ""
            for key in data_needed:
                if data_needed[key]:
                    elmnt += str(element[key]) + ','
            if elmnt[-1] == ',':
                elmnt = elmnt[:-1]
            elem_to_write.append(elmnt)
        file.write("\n".join(elem_to_write))
        file.write('\n')


def write_json(output, elements, data_needed):
    """
    Write selected data from elements to a JSON file.

    Parameters:
        output (str): The base name of the output file.
        elements (list): A list of dictionaries with element data.
        data_needed (dict): A dict indicating which keys to write.
    """
    with open(os.path.join(Path(__file__).parents[1], output + '.json'),
              'w', encoding="utf8") as file:
        elem_to_write = []

        for element in elements:
            elmnt = {}
            for key in data_needed:
                if data_needed[key]:
                    elmnt[key] = element[key]
            elem_to_write.append(elmnt)
        file.write(json.dumps(elem_to_write, indent=4))
        file.write('\n')


# pylint: disable=missing-function-docstring
def main():
    default_file = 'SpecificData'
    parser = create_commandeline_parser(default_file)

    args = parser.parse_args()
    elements, keys = read_periodic_table()

    data_needed = {}
    if args.properties:
        data_needed = parse_properties(data_needed, args, keys)
    if args.interactive:
        data_needed = parse_interactive(data_needed, keys)
    if data_needed:
        save2file(args, elements, data_needed, default_file)
    else:
        print(color_code.RED + 'No properties selected.' + color_code.END)


if __name__ == '__main__':
    main()