Script for organizing mineral data from Aurora 4x

mineral_extractor.py

# Mineral Extractor for Aurora 4x, by Steinar Lima (steili).
# - tested with Python 2.7, but most versions should be able to run it
#
# License shizzle: Do whatever you'd like with it! Feedback are welcomed at /r/aurora or IRC.
#
# Organizes and presents raw data from Aurora into a *.csv file. 
#
# ---- USAGE ----
#
# 1.  Create a folder next to this file named 'systems'
# 2.  Create an empty *.txt file for every system you have geological data for.
#     Use the system names as the file names
# 3.  For each system
#      a) locate it on the system map,
#      b) choose the 'Mineral' tab on the left hand side,
#      c) click the 'Mineral Text' button
#      d) copy all the text and paste it into the corresponding text file you created above
# 4.  When all this is done, just run the script, and a new file will be created named
#     "systems_mineral_data.csv" in the same folder as you ran the script from.
#     (double-clicking the script will run it if you have Python installed)
#
# ---- Import to Excel (one time thing, next time you can just click 'Refresh all' in the Excel 'Data' tab----
#
# 1.  Create a _new_ excel document
# 2.  Go to the "Data" tab, and click import "From text"
# 3.  Locate the "systems_mineral_data.csv"  created by this script, and click 'Open'
# 4.  Choose 'Delimited' and check the 'My data has headers' box -> click 'Next'
# 5.  Choose 'Comma' as the delimiter (remove any other delimiter options) -> click 'Next'
# 6.  Click 'Finish'
# 7.  Choose 'Add this data to the model', choose 'Table' and click 'Properties'.
# 8.  Uncheck the 'Prompt for file name on refresh', so you don't have to chose the
#     file every time you update the data
# 9.  Click OK and OK, the data should now appear in a table in Excel.
# 10. To avoid filling the table with junk as you update the table from time to time, click
#     'Properties', right next to the 'Update all' button in the 'Data' tab in Excel.
# 11. Choose 'Overwrite existing cells with new data, clear unused cells'
# 12. Remember that all changes yoy make to this table are overwritten
#     when you update the data (perhaps not the formatting)
# 10. Enjoy! :D
#
# For experienced Python programmers, look in the code below for some more
# advanced options like adding your own heuristics and adding regex patterns
# to clean up planet names (I for one prepend all my colonized planets of
# any importance with a some parentheses and information)

import os, glob, itertools, csv, math, re

# Regex pattern for modifying the body names. Leave blank if you don't want
# to modify anything. The resulting body is retrieved using group()[0] - make
# sure your match is in this group! (or just alter the code accordingly)
# Example:
#body_name_pattern = '(?:\(.+\))*(.+)'
#
# This will remove content inside a set of parantheses if it is located
# in the front of the string, e.g. "(C)Earth" would return "Earth"
body_name_pattern = ''

# Heuristics used for judging the quality of a mineral deposit.
# Each function should takes one argument (a list of tuples of
#  quantity and accessibility of each mineral, and should return the heuristic value
heuristics = [
    ('No. minerals', lambda l: len(l)),
    ('Tot. minerals', lambda l:sum(m[0] for m in l)),
    ('Sum. acc.', lambda l: sum(m[1] for m in l)),
    ('A^2 * log10(Q)^2', lambda l: sum(a**2 * math.log10(q**2) for q, a in l)),
]

mineral_headers = [
    'Duranium', 'Neutronium', 'Corbomite', 'Tritanium', 'Boronide',
    'Mercassium', 'Vendarite', 'Sorium', 'Uridium', 'Corundium', 'Gallicite',
]

# Suffix that are added to the mineral names for the accessibility columns
acc_suffix = '_Acc'

heurisitc_headers = [h[0] for h in heuristics]
headers = ['System', 'Body'] + heurisitc_headers

for name in mineral_headers:
    headers.append(name)
    headers.append('{}{}'.format(name, acc_suffix))

with open(os.path.join('systems_mineral_data.csv'), 'w') as out_file:
    writer = csv.DictWriter(
        out_file, headers, restval='', delimiter=',', lineterminator='\n'
    )
    writer.writeheader()

    # Walks over every *.txt file located in the "systems/" folder.
    # The system name is taken as the filename.
    for system in glob.glob('systems/*.txt'):
        system_name = system.replace('.txt', '').split('\\')[-1].title().strip()

        with open(system, 'rb') as system_file:
            # - Separate the groups by a line only containing a newline
            # - Filter out the newline groups by skipping every second group
            # - Drop the first group, we don't want the number of planets,
            #   surveys etc.
            system_bodies = itertools.islice(
                itertools.groupby(
                    system_file, lambda s: s == '\r\n'), 2, None, 2
            )

            # Loop over each system body, and save the system, name, and
            # quantity and accessibility of each mineral
            for _, body_minerals_data in system_bodies:
                raw_body_name = next(body_minerals_data).strip()

                # Alter body names if a regex pattern is given.
                body_name = raw_body_name if not body_name_pattern else re.match(
                                  body_name_pattern, raw_body_name).groups()[0].strip()

                # Create a dict that temporary stores the mineral data before it's written to file
                data = {'System': system_name, 'Body': body_name}
                heuristic_data = []

                # Loop over and save the quantity and acc. of each mineral
                # on the body
                for mineral_data in body_minerals_data:
                    mineral, quantity, acc = re.match(
                        '([\D]+)([\d\s]+)\sAcc: (.+)', mineral_data
                    ).groups()
                    q, a = int(quantity.strip().replace(' ', '')), float(acc.strip())
                    data[mineral.strip()] = q
                    data['{}{}'.format(mineral.strip(),acc_suffix)] = a
                    heuristic_data.append((q, a))

                # Calculate and sum all the heuristics for this mineral deposit
                for name, func in heuristics:
                    data[name] = func(heuristic_data)

                # Write the data to the output file
                writer.writerow(data)