<commit_before>'''
This script creates maps of model outputs at roughly 10km resolution (0.1x0.1 degrees), where each output pixel
represents the total value in the pixel (not the density) (hence, the aggregated results).
This is currently only set up for gross emissions from biomass+soil and net flux from biomass+soil.
It iterates through all the model outputs that are supplied.
First, it rewindows the model output, pixel area tile, and tcd tile into 400x400 (0.1x0.1 degree) windows, instead of the native
40000x1 pixel windows.
Then it calculates the per pixel value for each model output pixel and sums those values within each 0.1x0.1 degree
aggregated pixel.
It converts cumulative carbon gain to CO2 gain per year, converts cumulative CO2 flux to CO2 flux per year, and
converts cumulative gross CO2 emissions to gross CO2 emissions per year.
The user has to supply a tcd threshold for which forest pixels to include in the results.
'''

import multiprocessing
import aggregate_results_to_10_km
import subprocess
from functools import partial
import argparse
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main():

    # Files to download for this script. 'true'/'false' says whether the input directory and pattern should be
    # changed for a sensitivity analysis. This does not need to change based on what run is being done;
    # this assignment should be true for all sensitivity analyses and the standard model.
    download_dict = {
             # cn.gross_emis_all_gases_all_drivers_biomass_soil_dir: [cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil, 'true'],
             cn.cumul_gain_AGCO2_BGCO2_all_types_dir: [cn.pattern_cumul_gain_AGCO2_BGCO2_all_types],
             cn.net_flux_dir: [cn.pattern_net_flux]
             }

    # Sole argument for the script: the tree cover density threshold (pixels below this will not be aggregated)
    parser = argparse.ArgumentParser(description='Create maps of model output at approximately 10 km resolution')
    parser.add_argument('--tcd-threshold', '-tcd', required=True,
                        help='Tree cover density threshold above which pixels will be included in the aggregation.')
    parser.add_argument('--model-type', '-t', required=True,
                        help='{}'.format(cn.model_type_arg_help))
    parser.add_argument('--std-net-flux-aggreg', '-sagg', required=False,
                        help='The folder of the standard model net flux aggregated map, for comparison with the sensitivity analysis map')
    args = parser.parse_args()
    thresh = args.tcd_threshold
    thresh = int(thresh)
    sensit_type = args.model_type
    std_net_flux = args.std_net_flux_aggreg

    # Checks whether the canopy cover argument is valid
    if thresh < 0 or thresh > 99:
        raise Exception('Invalid tcd. Please provide an integer between 0 and 99.')

    # Checks whether the sensitivity analysis argument is valid
    uu.check_sensit_type(sensit_type)

    # tile_id_list = ['00N_110E'] # test tiles
    tile_id_list = 'all'

    # # Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
    # uu.s3_flexible_download(cn.pixel_area_dir, cn.pattern_pixel_area, '.', sensit_type, tile_id_list)
    # # tree cover density tiles-- necessary for filtering sums by tcd
    # uu.s3_flexible_download(cn.tcd_dir, cn.pattern_tcd, '.', sensit_type, tile_id_list)

    print "Model outputs to process are:", download_dict

    # List of output directories. Modified later for sensitivity analysis.
    # Output pattern is determined later.
    output_dir_list = [cn.output_aggreg_dir]

    # # Iterates through the types of tiles to be processed
    # for dir, download_pattern in download_dict.items():
    #
    #     download_pattern_name = download_pattern[0]
    #
    #     # Downloads the model output tiles to be processed
    #     uu.s3_flexible_download(dir, download_pattern_name, '.', sensit_type, tile_id_list)
    #
    #     # Renames the tiles according to the sensitivity analysis before creating dummy tiles.
    #     # The renaming function requires a whole tile name, so this passes a dummy time name that is then stripped a few
    #     # lines later.
    #     tile_id = 'XXXXXXXX'     # a dummy tile name. It is removed in the call to sensit_tile_rename
    #     output_pattern = uu.sensit_tile_rename(sensit_type, tile_id, download_pattern_name)
    #     pattern = output_pattern[9:-4]
    #
    #     # Lists the tiles of the particular type that is being iterates through.
    #     # Excludes all intermediate files
    #     tile_list = uu.tile_list_spot_machine(".", "{}.tif".format(pattern))
    #     # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
    #     tile_list = [i for i in tile_list if not ('hanson_2013' in i)]
    #     tile_list = [i for i in tile_list if not ('rewindow' in i)]
    #     tile_list = [i for i in tile_list if not ('10km' in i)]
    #
    #     print tile_list
    #     print "There are {} tiles to process".format(str(len(tile_list))) + "\n"
    #     print "Processing:", dir, "; ", pattern
    #
    #     # Converts the 10x10 degree Hansen tiles that are in windows of 40000x1 pixels to windows of 400x400 pixels,
    #     # which is the resolution of the output tiles. This will allow the 30x30 m pixels in each window to be summed.
    #     # For multiprocessor use. count/2 used about 400 GB of memory on an r4.16xlarge machine, so that was okay.
    #     count = multiprocessing.cpu_count()
    #     pool = multiprocessing.Pool(count/2)
    #     pool.map(aggregate_results_to_10_km.rewindow, tile_list)
    #     # Added these in response to error12: Cannot allocate memory error.
    #     # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    #     # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    #     pool.close()
    #     pool.join()
    #
    #     # # For single processor use
    #     # for tile in tile_list:
    #     #
    #     #     aggregate_results_to_10_km.rewindow(tile)
    #
    #     # Converts the existing (per ha) values to per pixel values (e.g., emissions/ha to emissions/pixel)
    #     # and sums those values in each 400x400 pixel window.
    #     # The sum for each 400x400 pixel window is stored in a 2D array, which is then converted back into a raster at
    #     # 0.1x0.1 degree resolution (approximately 10m in the tropics).
    #     # Each pixel in that raster is the sum of the 30m pixels converted to value/pixel (instead of value/ha).
    #     # The 0.1x0.1 degree tile is output.
    #     # For multiprocessor use. This used about 450 GB of memory with count/2, it's okay on an r4.16xlarge
    #     count = multiprocessing.cpu_count()
    #     pool = multiprocessing.Pool(count/2)
    #     pool.map(partial(aggregate_results_to_10_km.aggregate, thresh=thresh), tile_list)
    #     # Added these in response to error12: Cannot allocate memory error.
    #     # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    #     # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    #     pool.close()
    #     pool.join()
    #
    #     # # For single processor use
    #     # for tile in tile_list:
    #     #
    #     #     aggregate_results_to_10_km.aggregate(tile, thresh)
    #
    #     # Makes a vrt of all the output 10x10 tiles (10 km resolution)
    #     out_vrt = "{}_10km.vrt".format(pattern)
    #     os.system('gdalbuildvrt -tr 0.1 0.1 {0} *{1}_10km*.tif'.format(out_vrt, pattern))
    #
    #     # Creates the output name for the 10km map
    #     out_pattern = uu.name_aggregated_output(download_pattern_name, thresh, sensit_type)
    #
    #     # Produces a single raster of all the 10x10 tiles (10 km resolution)
    #     cmd = ['gdalwarp', '-t_srs', "EPSG:4326", '-overwrite', '-dstnodata', '0', '-co', 'COMPRESS=LZW',
    #            '-tr', '0.1', '0.1',
    #            out_vrt, '{}.tif'.format(out_pattern)]
    #     subprocess.check_call(cmd)
    #
    #     print "Tiles processed. Uploading to s3 now..."
    #
    #     # # Cleans up the folder before starting on the next raster type
    #     # vrtList = glob.glob('*vrt')
    #     # for vrt in vrtList:
    #     #     os.remove(vrt)
    #     #
    #     # for tile_id in tile_list:
    #     #     os.remove('{0}_{1}.tif'.format(tile_id, pattern))
    #     #     os.remove('{0}_{1}_rewindow.tif'.format(tile_id, pattern))
    #     #     os.remove('{0}_{1}_10km.tif'.format(tile_id, pattern))
    #
    #     # If the model run isn't the standard one, the output directory is changed
    #     if sensit_type != 'std':
    #         print "Changing output directory and file name pattern based on sensitivity analysis"
    #         output_dir_list = uu.alter_dirs(sensit_type, output_dir_list)
    #
    #     # Uploads all output tiles to s3
    #     uu.upload_final_set(output_dir_list[0], out_pattern)

    if sensit_type != 'std':

        # Copies the standard model aggregation outputs to s3
        uu.s3_flexible_download(std_net_flux, cn.pattern_aggreg, '.', 'std', 'all')

        std_aggreg_flux = glob.glob('net_flux*std*')[0]
        sensit_aggreg_flux = glob.glob('net_flux*{}*'.format(sensit_type))[0]

        print "Standard model net flux:", std_aggreg_flux
        print "Sensitivity model net flux:", sensit_aggreg_flux

        print "Creating map of percent difference between standard and {} net flux".format(sensit_type)
        aggregate_results_to_10_km.percent_diff(std_aggreg_flux, sensit_aggreg_flux, sensit_type)

        print output_dir_list
        print cn.pattern_aggreg_perc_diff
        # uu.upload_final_set(output_dir_list[0], cn.pattern_aggreg_perc_diff)

if __name__ == '__main__':
    main()<commit_msg>Working on modifying 10km aggregate output script to calculate percent different between aggregated standard net flux and sensitivity model net flux.<commit_after>'''
This script creates maps of model outputs at roughly 10km resolution (0.1x0.1 degrees), where each output pixel
represents the total value in the pixel (not the density) (hence, the aggregated results).
This is currently only set up for gross emissions from biomass+soil and net flux from biomass+soil.
It iterates through all the model outputs that are supplied.
First, it rewindows the model output, pixel area tile, and tcd tile into 400x400 (0.1x0.1 degree) windows, instead of the native
40000x1 pixel windows.
Then it calculates the per pixel value for each model output pixel and sums those values within each 0.1x0.1 degree
aggregated pixel.
It converts cumulative carbon gain to CO2 gain per year, converts cumulative CO2 flux to CO2 flux per year, and
converts cumulative gross CO2 emissions to gross CO2 emissions per year.
The user has to supply a tcd threshold for which forest pixels to include in the results.
'''

import multiprocessing
import aggregate_results_to_10_km
import subprocess
from functools import partial
import argparse
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

def main():

    # Files to download for this script. 'true'/'false' says whether the input directory and pattern should be
    # changed for a sensitivity analysis. This does not need to change based on what run is being done;
    # this assignment should be true for all sensitivity analyses and the standard model.
    download_dict = {
             # cn.gross_emis_all_gases_all_drivers_biomass_soil_dir: [cn.pattern_gross_emis_all_gases_all_drivers_biomass_soil, 'true'],
             cn.cumul_gain_AGCO2_BGCO2_all_types_dir: [cn.pattern_cumul_gain_AGCO2_BGCO2_all_types],
             cn.net_flux_dir: [cn.pattern_net_flux]
             }

    # Sole argument for the script: the tree cover density threshold (pixels below this will not be aggregated)
    parser = argparse.ArgumentParser(description='Create maps of model output at approximately 10 km resolution')
    parser.add_argument('--tcd-threshold', '-tcd', required=True,
                        help='Tree cover density threshold above which pixels will be included in the aggregation.')
    parser.add_argument('--model-type', '-t', required=True,
                        help='{}'.format(cn.model_type_arg_help))
    parser.add_argument('--std-net-flux-aggreg', '-sagg', required=False,
                        help='The folder of the standard model net flux aggregated map, for comparison with the sensitivity analysis map')
    args = parser.parse_args()
    thresh = args.tcd_threshold
    thresh = int(thresh)
    sensit_type = args.model_type
    std_net_flux = args.std_net_flux_aggreg

    # Checks whether the canopy cover argument is valid
    if thresh < 0 or thresh > 99:
        raise Exception('Invalid tcd. Please provide an integer between 0 and 99.')

    # Checks whether the sensitivity analysis argument is valid
    uu.check_sensit_type(sensit_type)

    # tile_id_list = ['00N_110E'] # test tiles
    tile_id_list = 'all'

    # # Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
    # uu.s3_flexible_download(cn.pixel_area_dir, cn.pattern_pixel_area, '.', sensit_type, tile_id_list)
    # # tree cover density tiles-- necessary for filtering sums by tcd
    # uu.s3_flexible_download(cn.tcd_dir, cn.pattern_tcd, '.', sensit_type, tile_id_list)

    print "Model outputs to process are:", download_dict

    # List of output directories. Modified later for sensitivity analysis.
    # Output pattern is determined later.
    output_dir_list = [cn.output_aggreg_dir]

    # If the model run isn't the standard one, the output directory is changed
    if sensit_type != 'std':
        print "Changing output directory and file name pattern based on sensitivity analysis"
        output_dir_list = uu.alter_dirs(sensit_type, output_dir_list)

    # # Iterates through the types of tiles to be processed
    # for dir, download_pattern in download_dict.items():
    #
    #     download_pattern_name = download_pattern[0]
    #
    #     # Downloads the model output tiles to be processed
    #     uu.s3_flexible_download(dir, download_pattern_name, '.', sensit_type, tile_id_list)
    #
    #     # Renames the tiles according to the sensitivity analysis before creating dummy tiles.
    #     # The renaming function requires a whole tile name, so this passes a dummy time name that is then stripped a few
    #     # lines later.
    #     tile_id = 'XXXXXXXX'     # a dummy tile name. It is removed in the call to sensit_tile_rename
    #     output_pattern = uu.sensit_tile_rename(sensit_type, tile_id, download_pattern_name)
    #     pattern = output_pattern[9:-4]
    #
    #     # Lists the tiles of the particular type that is being iterates through.
    #     # Excludes all intermediate files
    #     tile_list = uu.tile_list_spot_machine(".", "{}.tif".format(pattern))
    #     # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
    #     tile_list = [i for i in tile_list if not ('hanson_2013' in i)]
    #     tile_list = [i for i in tile_list if not ('rewindow' in i)]
    #     tile_list = [i for i in tile_list if not ('10km' in i)]
    #
    #     print tile_list
    #     print "There are {} tiles to process".format(str(len(tile_list))) + "\n"
    #     print "Processing:", dir, "; ", pattern
    #
    #     # Converts the 10x10 degree Hansen tiles that are in windows of 40000x1 pixels to windows of 400x400 pixels,
    #     # which is the resolution of the output tiles. This will allow the 30x30 m pixels in each window to be summed.
    #     # For multiprocessor use. count/2 used about 400 GB of memory on an r4.16xlarge machine, so that was okay.
    #     count = multiprocessing.cpu_count()
    #     pool = multiprocessing.Pool(count/2)
    #     pool.map(aggregate_results_to_10_km.rewindow, tile_list)
    #     # Added these in response to error12: Cannot allocate memory error.
    #     # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    #     # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    #     pool.close()
    #     pool.join()
    #
    #     # # For single processor use
    #     # for tile in tile_list:
    #     #
    #     #     aggregate_results_to_10_km.rewindow(tile)
    #
    #     # Converts the existing (per ha) values to per pixel values (e.g., emissions/ha to emissions/pixel)
    #     # and sums those values in each 400x400 pixel window.
    #     # The sum for each 400x400 pixel window is stored in a 2D array, which is then converted back into a raster at
    #     # 0.1x0.1 degree resolution (approximately 10m in the tropics).
    #     # Each pixel in that raster is the sum of the 30m pixels converted to value/pixel (instead of value/ha).
    #     # The 0.1x0.1 degree tile is output.
    #     # For multiprocessor use. This used about 450 GB of memory with count/2, it's okay on an r4.16xlarge
    #     count = multiprocessing.cpu_count()
    #     pool = multiprocessing.Pool(count/2)
    #     pool.map(partial(aggregate_results_to_10_km.aggregate, thresh=thresh), tile_list)
    #     # Added these in response to error12: Cannot allocate memory error.
    #     # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    #     # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    #     pool.close()
    #     pool.join()
    #
    #     # # For single processor use
    #     # for tile in tile_list:
    #     #
    #     #     aggregate_results_to_10_km.aggregate(tile, thresh)
    #
    #     # Makes a vrt of all the output 10x10 tiles (10 km resolution)
    #     out_vrt = "{}_10km.vrt".format(pattern)
    #     os.system('gdalbuildvrt -tr 0.1 0.1 {0} *{1}_10km*.tif'.format(out_vrt, pattern))
    #
    #     # Creates the output name for the 10km map
    #     out_pattern = uu.name_aggregated_output(download_pattern_name, thresh, sensit_type)
    #
    #     # Produces a single raster of all the 10x10 tiles (10 km resolution)
    #     cmd = ['gdalwarp', '-t_srs', "EPSG:4326", '-overwrite', '-dstnodata', '0', '-co', 'COMPRESS=LZW',
    #            '-tr', '0.1', '0.1',
    #            out_vrt, '{}.tif'.format(out_pattern)]
    #     subprocess.check_call(cmd)
    #
    #     print "Tiles processed. Uploading to s3 now..."
    #
    #     # # Cleans up the folder before starting on the next raster type
    #     # vrtList = glob.glob('*vrt')
    #     # for vrt in vrtList:
    #     #     os.remove(vrt)
    #     #
    #     # for tile_id in tile_list:
    #     #     os.remove('{0}_{1}.tif'.format(tile_id, pattern))
    #     #     os.remove('{0}_{1}_rewindow.tif'.format(tile_id, pattern))
    #     #     os.remove('{0}_{1}_10km.tif'.format(tile_id, pattern))
    #
    #     # Uploads all output tiles to s3
    #     uu.upload_final_set(output_dir_list[0], out_pattern)

    if sensit_type != 'std':

        # Copies the standard model aggregation outputs to s3
        uu.s3_flexible_download(std_net_flux, cn.pattern_aggreg, '.', 'std', 'all')

        std_aggreg_flux = glob.glob('net_flux*std*')[0]
        sensit_aggreg_flux = glob.glob('net_flux*{}*'.format(sensit_type))[0]

        print "Standard model net flux:", std_aggreg_flux
        print "Sensitivity model net flux:", sensit_aggreg_flux

        print "Creating map of percent difference between standard and {} net flux".format(sensit_type)
        aggregate_results_to_10_km.percent_diff(std_aggreg_flux, sensit_aggreg_flux, sensit_type)

        print output_dir_list
        print cn.pattern_aggreg_perc_diff
        # uu.upload_final_set(output_dir_list[0], cn.pattern_aggreg_perc_diff)

if __name__ == '__main__':
    main()<|endoftext|>


<commit_before>import ttag

from django import template
from django.template.loader import render_to_string
from django.db.models.query import QuerySet
from django.conf import settings
from django.contrib.gis.geos import Point
from django.contrib.gis.maps.google.zoom import GoogleZoom

from geotagging.models import PointGeoTag

register = template.Library()

class Javascript(ttag.Tag):

    class Meta:
        block = False
        name = 'geotagging_maps_api'

    def render(self, context):
        t = template.loader.get_template('geotagging/map_scripts.html')
        return t.render(template.Context({}))

register.tag(Javascript)

"""
Loosly based on dajngo-easy-maps: https://bitbucket.org/kmike/django-easy-maps/overview
"""
class MapObjects(ttag.Tag):
    class Meta:
        name = 'geotagging_map'

    objects = ttag.Arg(positional=True)
    width = ttag.Arg(required=False)
    height = ttag.Arg(required=False)
    zoom = ttag.Arg(required=False)
    static = ttag.Arg(required=False)
    extra = ttag.Arg(required=False)

    def render(self, context):
        data = self.resolve(context)
        objects = data.get('objects', None)
        width = data.get('width', None)
        height = data.get('height', None)
        zoom = data.get('zoom', None)
        static = data.get('static', None) == "true"
        extra = data.get('extra', None)

        context['request'].session['geotagging_map_counter'] = (
            context['request'].session.get('geotagging_map_counter', 0) + 1)
        id = context['request'].session['geotagging_map_counter']

        if isinstance(objects, PointGeoTag):
            latlng = objects.get_point_coordinates(as_string=True, inverted=True)
            markers = [{'latlng':latlng, 'object': objects}]
        elif isinstance(objects, basestring):
            latlng = objects
            markers = [{'latlng':latlng}]
        elif isinstance(objects, QuerySet) or isinstance(objects, list):
            if len(objects) == 0:
                centroid = None
            elif not getattr(settings, 'USE_GEOGRAPHY', True):
                centroid = objects.collect().envelope.centroid
                gz = GoogleZoom()
                zoom = gz.get_zoom(objects.unionagg())
            else:
                if static:
                    centroid = None
                else:
                    centroid = objects[0].get_point_coordinates(as_string=False,
                                                                inverted=True)
            if centroid:
                latlng = '%s,%s' % (centroid[0], centroid[1])
            else:
                latlng = None
            markers = [{'latlng': i.get_point_coordinates(as_string=True, inverted=True),
                        'object': i} for i in objects]
        else:
            raise template.TemplateSyntaxError(
                'The first parameter must be either a PointGeoTag subclass, '
                'a queryset of PointGeoTag subclasses, '
                'a list of PointGeoTag subclases or a LatLong string. '
                'A %s was given' % type(objects))

        if static:
            t = template.loader.get_template('geotagging/staticmap.html')
            return t.render(template.Context({
                        'map_id': id,
                        'width': width,
                        'height': height,
                        'latlng': latlng,
                        'markers': markers,
                        'extra': extra,
                        'zoom': zoom,
                        }))

        t = template.loader.get_template('geotagging/map.html')
        return t.render(template.Context({
                    'map_id': id,
                    'width': width,
                    'height': height,
                    'latlng': latlng,
                    'markers': markers,
                    'extra': extra,
                    'zoom': zoom,
                    }))

register.tag(MapObjects)

"""
Notes:

Static:
 {% include_map_js %} - Adds the necessary javascript for the maps to load
 {% map_objects queryset %} - Creates a map showing all objects in the queryset
Interactive:
 {% map_objects Model %} - Creates a map showing all Model objects within the map section
 {% map_objects Model filter %} - Same as above, but adds filtering based on a function

Queue:
 ???
 {% include_queue_js element_id %} - Inncludes the javascript for the element queue
 {% show_queue %}

"""
<commit_msg>updating ttags to latest version<commit_after>import ttag

from django import template
from django.template.loader import render_to_string
from django.db.models.query import QuerySet
from django.conf import settings
from django.contrib.gis.geos import Point
from django.contrib.gis.maps.google.zoom import GoogleZoom

from geotagging.models import PointGeoTag

register = template.Library()

class Javascript(ttag.Tag):

    class Meta:
        block = False
        name = 'geotagging_maps_api'

    def render(self, context):
        t = template.loader.get_template('geotagging/map_scripts.html')
        return t.render(template.Context({}))

register.tag(Javascript)

"""
Loosly based on dajngo-easy-maps: https://bitbucket.org/kmike/django-easy-maps/overview
"""
class MapObjects(ttag.Tag):
    class Meta:
        name = 'geotagging_map'

    objects = ttag.Arg(positional=True)
    width = ttag.Arg(required=False, keyword=True)
    height = ttag.Arg(required=False, keyword=True)
    zoom = ttag.Arg(required=False, keyword=True)
    static = ttag.Arg(required=False, keyword=True)
    extra = ttag.Arg(required=False, keyword=True)

    def render(self, context):
        data = self.resolve(context)
        objects = data.get('objects', None)
        width = data.get('width', None)
        height = data.get('height', None)
        zoom = data.get('zoom', None)
        static = data.get('static', None) == "true"
        extra = data.get('extra', None)

        context['request'].session['geotagging_map_counter'] = (
            context['request'].session.get('geotagging_map_counter', 0) + 1)
        id = context['request'].session['geotagging_map_counter']

        if isinstance(objects, PointGeoTag):
            latlng = objects.get_point_coordinates(as_string=True, inverted=True)
            markers = [{'latlng':latlng, 'object': objects}]
        elif isinstance(objects, basestring):
            latlng = objects
            markers = [{'latlng':latlng}]
        elif isinstance(objects, QuerySet) or isinstance(objects, list):
            if len(objects) == 0:
                centroid = None
            elif not getattr(settings, 'USE_GEOGRAPHY', True):
                centroid = objects.collect().envelope.centroid
                gz = GoogleZoom()
                zoom = gz.get_zoom(objects.unionagg())
            else:
                if static:
                    centroid = None
                else:
                    centroid = objects[0].get_point_coordinates(as_string=False,
                                                                inverted=True)
            if centroid:
                latlng = '%s,%s' % (centroid[0], centroid[1])
            else:
                latlng = None
            markers = [{'latlng': i.get_point_coordinates(as_string=True, inverted=True),
                        'object': i} for i in objects]
        else:
            raise template.TemplateSyntaxError(
                'The first parameter must be either a PointGeoTag subclass, '
                'a queryset of PointGeoTag subclasses, '
                'a list of PointGeoTag subclases or a LatLong string. '
                'A %s was given' % type(objects))

        if static:
            t = template.loader.get_template('geotagging/staticmap.html')
            return t.render(template.Context({
                        'map_id': id,
                        'width': width,
                        'height': height,
                        'latlng': latlng,
                        'markers': markers,
                        'extra': extra,
                        'zoom': zoom,
                        }))

        t = template.loader.get_template('geotagging/map.html')
        return t.render(template.Context({
                    'map_id': id,
                    'width': width,
                    'height': height,
                    'latlng': latlng,
                    'markers': markers,
                    'extra': extra,
                    'zoom': zoom,
                    }))

register.tag(MapObjects)

"""
Notes:

Static:
 {% include_map_js %} - Adds the necessary javascript for the maps to load
 {% map_objects queryset %} - Creates a map showing all objects in the queryset
Interactive:
 {% map_objects Model %} - Creates a map showing all Model objects within the map section
 {% map_objects Model filter %} - Same as above, but adds filtering based on a function

Queue:
 ???
 {% include_queue_js element_id %} - Inncludes the javascript for the element queue
 {% show_queue %}

"""
<|endoftext|>


<commit_before># Copyright 2014 Open Source Robotics Foundation, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This module provides tools for colorizing terminal output.

This module defines the ansi escape sequences used for colorizing the
output from terminal programs in Linux.
You can access the ansi escape sequences using the :py:func:`ansi` function:

.. code-block:: python

    >>> from osrf_pycommon.terminal_color import ansi
    >>> print(["This is ", ansi('red'), "red", ansi('reset'), "."])
    ['This is ', '\\x1b[31m', 'red', '\\x1b[0m', '.']

You can also use :py:func:`format_color` to do in-line substitution of keys
wrapped in ``@{}`` markers for their ansi escape sequences:

.. code-block:: python

    >>> from osrf_pycommon.terminal_color import format_color
    >>> print(format_color("This is @{bf}blue@{reset}.").split())
    ['This', 'is', '\\x1b[34mblue\\x1b[0m.']

This is a list of all of the available substitutions:

+-------------------+--------------+--------------+
| Long Form         | Shorter      | Value        |
+===================+==============+==============+
| ``@{blackf}``     | ``@{kf}``    | ``\\033[30m`` |
+-------------------+--------------+--------------+
| ``@{redf}``       | ``@{rf}``    | ``\\033[31m`` |
+-------------------+--------------+--------------+
| ``@{greenf}``     | ``@{gf}``    | ``\\033[32m`` |
+-------------------+--------------+--------------+
| ``@{yellowf}``    | ``@{yf}``    | ``\\033[33m`` |
+-------------------+--------------+--------------+
| ``@{bluef}``      | ``@{bf}``    | ``\\033[34m`` |
+-------------------+--------------+--------------+
| ``@{purplef}``    | ``@{pf}``    | ``\\033[35m`` |
+-------------------+--------------+--------------+
| ``@{cyanf}``      | ``@{cf}``    | ``\\033[36m`` |
+-------------------+--------------+--------------+
| ``@{whitef}``     | ``@{wf}``    | ``\\033[37m`` |
+-------------------+--------------+--------------+
| ``@{blackb}``     | ``@{kb}``    | ``\\033[40m`` |
+-------------------+--------------+--------------+
| ``@{redb}``       | ``@{rb}``    | ``\\033[41m`` |
+-------------------+--------------+--------------+
| ``@{greenb}``     | ``@{gb}``    | ``\\033[42m`` |
+-------------------+--------------+--------------+
| ``@{yellowb}``    | ``@{yb}``    | ``\\033[43m`` |
+-------------------+--------------+--------------+
| ``@{blueb}``      | ``@{bb}``    | ``\\033[44m`` |
+-------------------+--------------+--------------+
| ``@{purpleb}``    | ``@{pb}``    | ``\\033[45m`` |
+-------------------+--------------+--------------+
| ``@{cyanb}``      | ``@{cb}``    | ``\\033[46m`` |
+-------------------+--------------+--------------+
| ``@{whiteb}``     | ``@{wb}``    | ``\\033[47m`` |
+-------------------+--------------+--------------+
| ``@{escape}``     |              | ``\\033``     |
+-------------------+--------------+--------------+
| ``@{reset}``      | ``@|``       | ``\\033[0m``  |
+-------------------+--------------+--------------+
| ``@{boldon}``     | ``@!``       | ``\\033[1m``  |
+-------------------+--------------+--------------+
| ``@{italicson}``  | ``@/``       | ``\\033[3m``  |
+-------------------+--------------+--------------+
| ``@{ulon}``       | ``@_``       | ``\\033[4m``  |
+-------------------+--------------+--------------+
| ``@{invon}``      |              | ``\\033[7m``  |
+-------------------+--------------+--------------+
| ``@{boldoff}``    |              | ``\\033[22m`` |
+-------------------+--------------+--------------+
| ``@{italicsoff}`` |              | ``\\033[23m`` |
+-------------------+--------------+--------------+
| ``@{uloff}``      |              | ``\\033[24m`` |
+-------------------+--------------+--------------+
| ``@{invoff}``     |              | ``\\033[27m`` |
+-------------------+--------------+--------------+

These substitution's values come from the ANSI color escape sequences,
see: http://en.wikipedia.org/wiki/ANSI_escape_code

Also for any of the keys which have a trailing ``f``, you can safely drop the
trailing ``f`` and get the same thing.

For example, ``format_color("@{redf}")`` and ``format_color("@{red}")`` are
functionally equivalent.

Also, many of the substitutions have shorten forms for convenience, such that
``@{redf}``, ``@{rf}``, ``@{red}``, and ``@{r}`` are all the same.

Note that a trailing ``b`` is always required when specifying a background.

Some of the most common non-color sequences have ``{}``'less versions.

For example, ``@{boldon}``'s shorter form is ``@!``.

By default, the substitutions (and calls to :py:func:`ansi`) resolve to escape
sequences, but if you call :py:func:`disable_ansi_color_substitution_globally`
then they will resolve to empty strings.

This allows you to always use the substitution strings and disable them
globally when desired.

On Windows the substitutions are always resolved to empty strings as the
ansi escape sequences do not work on Windows.
Instead strings annotated with ``@{}`` style substitutions or raw ``\\x1b[xxm``
style ansi escape sequences must be passed to :py:func:`print_color` in order
for colors to be displayed on windows.
Also the :py:func:`print_ansi_color_win32` function can be used on strings
which only contain ansi escape sequences.
"""

from .ansi_re import remove_ansi_escape_senquences
from .ansi_re import split_by_ansi_escape_sequence

from .impl import ansi
from .impl import disable_ansi_color_substitution_globally
from .impl import enable_ansi_color_substitution_globally
from .impl import format_color
from .impl import get_ansi_dict
from .impl import print_ansi_color_win32
from .impl import print_color
from .impl import sanitize
from .impl import test_colors

__all__ = [
    'ansi',
    'disable_ansi_color_substitution_globally',
    'enable_ansi_color_substitution_globally',
    'format_color',
    'get_ansi_dict',
    'print_ansi_color_win32',
    'print_color',
    'remove_ansi_escape_senquences',
    'sanitize',
    'split_by_ansi_escape_sequence',
    'test_colors',
]
<commit_msg>terminal_color: add doc note about why external modules were not used<commit_after># Copyright 2014 Open Source Robotics Foundation, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
This module provides tools for colorizing terminal output.

This module defines the ansi escape sequences used for colorizing the
output from terminal programs in Linux.
You can access the ansi escape sequences using the :py:func:`ansi` function:

.. code-block:: python

    >>> from osrf_pycommon.terminal_color import ansi
    >>> print(["This is ", ansi('red'), "red", ansi('reset'), "."])
    ['This is ', '\\x1b[31m', 'red', '\\x1b[0m', '.']

You can also use :py:func:`format_color` to do in-line substitution of keys
wrapped in ``@{}`` markers for their ansi escape sequences:

.. code-block:: python

    >>> from osrf_pycommon.terminal_color import format_color
    >>> print(format_color("This is @{bf}blue@{reset}.").split())
    ['This', 'is', '\\x1b[34mblue\\x1b[0m.']

This is a list of all of the available substitutions:

+-------------------+--------------+--------------+
| Long Form         | Shorter      | Value        |
+===================+==============+==============+
| ``@{blackf}``     | ``@{kf}``    | ``\\033[30m`` |
+-------------------+--------------+--------------+
| ``@{redf}``       | ``@{rf}``    | ``\\033[31m`` |
+-------------------+--------------+--------------+
| ``@{greenf}``     | ``@{gf}``    | ``\\033[32m`` |
+-------------------+--------------+--------------+
| ``@{yellowf}``    | ``@{yf}``    | ``\\033[33m`` |
+-------------------+--------------+--------------+
| ``@{bluef}``      | ``@{bf}``    | ``\\033[34m`` |
+-------------------+--------------+--------------+
| ``@{purplef}``    | ``@{pf}``    | ``\\033[35m`` |
+-------------------+--------------+--------------+
| ``@{cyanf}``      | ``@{cf}``    | ``\\033[36m`` |
+-------------------+--------------+--------------+
| ``@{whitef}``     | ``@{wf}``    | ``\\033[37m`` |
+-------------------+--------------+--------------+
| ``@{blackb}``     | ``@{kb}``    | ``\\033[40m`` |
+-------------------+--------------+--------------+
| ``@{redb}``       | ``@{rb}``    | ``\\033[41m`` |
+-------------------+--------------+--------------+
| ``@{greenb}``     | ``@{gb}``    | ``\\033[42m`` |
+-------------------+--------------+--------------+
| ``@{yellowb}``    | ``@{yb}``    | ``\\033[43m`` |
+-------------------+--------------+--------------+
| ``@{blueb}``      | ``@{bb}``    | ``\\033[44m`` |
+-------------------+--------------+--------------+
| ``@{purpleb}``    | ``@{pb}``    | ``\\033[45m`` |
+-------------------+--------------+--------------+
| ``@{cyanb}``      | ``@{cb}``    | ``\\033[46m`` |
+-------------------+--------------+--------------+
| ``@{whiteb}``     | ``@{wb}``    | ``\\033[47m`` |
+-------------------+--------------+--------------+
| ``@{escape}``     |              | ``\\033``     |
+-------------------+--------------+--------------+
| ``@{reset}``      | ``@|``       | ``\\033[0m``  |
+-------------------+--------------+--------------+
| ``@{boldon}``     | ``@!``       | ``\\033[1m``  |
+-------------------+--------------+--------------+
| ``@{italicson}``  | ``@/``       | ``\\033[3m``  |
+-------------------+--------------+--------------+
| ``@{ulon}``       | ``@_``       | ``\\033[4m``  |
+-------------------+--------------+--------------+
| ``@{invon}``      |              | ``\\033[7m``  |
+-------------------+--------------+--------------+
| ``@{boldoff}``    |              | ``\\033[22m`` |
+-------------------+--------------+--------------+
| ``@{italicsoff}`` |              | ``\\033[23m`` |
+-------------------+--------------+--------------+
| ``@{uloff}``      |              | ``\\033[24m`` |
+-------------------+--------------+--------------+
| ``@{invoff}``     |              | ``\\033[27m`` |
+-------------------+--------------+--------------+

These substitution's values come from the ANSI color escape sequences,
see: http://en.wikipedia.org/wiki/ANSI_escape_code

Also for any of the keys which have a trailing ``f``, you can safely drop the
trailing ``f`` and get the same thing.

For example, ``format_color("@{redf}")`` and ``format_color("@{red}")`` are
functionally equivalent.

Also, many of the substitutions have shorten forms for convenience, such that
``@{redf}``, ``@{rf}``, ``@{red}``, and ``@{r}`` are all the same.

Note that a trailing ``b`` is always required when specifying a background.

Some of the most common non-color sequences have ``{}``'less versions.

For example, ``@{boldon}``'s shorter form is ``@!``.

By default, the substitutions (and calls to :py:func:`ansi`) resolve to escape
sequences, but if you call :py:func:`disable_ansi_color_substitution_globally`
then they will resolve to empty strings.

This allows you to always use the substitution strings and disable them
globally when desired.

On Windows the substitutions are always resolved to empty strings as the
ansi escape sequences do not work on Windows.
Instead strings annotated with ``@{}`` style substitutions or raw ``\\x1b[xxm``
style ansi escape sequences must be passed to :py:func:`print_color` in order
for colors to be displayed on windows.
Also the :py:func:`print_ansi_color_win32` function can be used on strings
which only contain ansi escape sequences.

.. note::

    There are existing Python modules like
    `colorama <https://pypi.python.org/pypi/colorama>`_ which provide ansi
    colorization on multiple platforms, so a valid question is:
    "why write this module?".
    The reason for writing this module is to provide the color annotation of
    strings and functions for removing or replacing ansi escape sequences which
    are not provided by modules like colorama.
    This module could have depended on colorama for colorization on Windows,
    but colorama works by replacing the built-in ``sys.stdout`` and
    ``sys.stderr``, which we did not want and it has extra functionality that
    we do not need.
    So, instead of depending on colorama, the Windows color printing code
    was used as the inspiration for the Windows color printing in the
    ``windows.py`` module in this ``terminal_color`` package.
    The colorama license was placed in the header of that file and the colorama
    license is compatible with this package's license.

"""

from .ansi_re import remove_ansi_escape_senquences
from .ansi_re import split_by_ansi_escape_sequence

from .impl import ansi
from .impl import disable_ansi_color_substitution_globally
from .impl import enable_ansi_color_substitution_globally
from .impl import format_color
from .impl import get_ansi_dict
from .impl import print_ansi_color_win32
from .impl import print_color
from .impl import sanitize
from .impl import test_colors

__all__ = [
    'ansi',
    'disable_ansi_color_substitution_globally',
    'enable_ansi_color_substitution_globally',
    'format_color',
    'get_ansi_dict',
    'print_ansi_color_win32',
    'print_color',
    'remove_ansi_escape_senquences',
    'sanitize',
    'split_by_ansi_escape_sequence',
    'test_colors',
]
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""Bootstrap the setup tool.

This assumes Python is installed on the target os, but not specifically Python3.

What this does (only intended for Mac atm):

* installs Homebrew
* Homebrew installs a core set of packages (git and python3)
* git check out the project into ~/setup
* run
  - setup (will restart os functions to reflect new settings)
  - setup brew
  - setup packages
* tell the user to restart terminal to get new everything

You should be able to run this with curl | python shenanigans.

"""

import os
import subprocess

REPO_URL = 'https://github.com/kbd/setup.git'
HOMEBREW_INSTALL_CMD = '/usr/bin/ruby -e "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install)"'

def main():
    print("Installing Homebrew")
    if not subprocess.call(['which', 'brew']):
        print("Homebrew is installed")
    else:
        subprocess.check_call(HOMEBREW_INSTALL_CMD, shell=True, executable='/bin/bash')

    print("Installing git and Python 3")
    subprocess.check_call("brew install git || brew upgrade git", shell=True)
    subprocess.check_call("brew install python || brew upgrade python", shell=True)
    subprocess.check_call(['pip3', 'install', '--upgrade', 'click'])  # required for 'setup'

    setup_path = os.path.expanduser('~/setup')
    if os.path.exists(setup_path):
        print("Setup location already exists, updating")
        subprocess.check_call(['git', 'pull'], cwd=setup_path)
    else:
        print("Checking out setup repo")
        subprocess.check_call(['git', 'clone', REPO_URL], cwd=os.path.dirname(setup_path))

    setup_exe = os.path.join(setup_path, 'HOME/bin/setup')

    print("Installing all the things")
    subprocess.check_call([setup_exe, 'init'])
    print("Done installing all the things. Restart your terminal.")

if __name__ == '__main__':
    main()
<commit_msg>Update Homebrew install command<commit_after>#!/usr/bin/env python

"""Bootstrap the setup tool.

This assumes Python is installed on the target os, but not specifically Python3.

What this does (only intended for Mac atm):

* installs Homebrew
* Homebrew installs a core set of packages (git and python3)
* git check out the project into ~/setup
* run
  - setup (will restart os functions to reflect new settings)
  - setup brew
  - setup packages
* tell the user to restart terminal to get new everything

You should be able to run this with curl | python shenanigans.

"""

import os
import subprocess

REPO_URL = 'https://github.com/kbd/setup.git'
HOMEBREW_INSTALL_CMD = '/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"'

def main():
    print("Installing Homebrew")
    if not subprocess.call(['which', 'brew']):
        print("Homebrew is installed")
    else:
        subprocess.check_call(HOMEBREW_INSTALL_CMD, shell=True, executable='/bin/bash')

    print("Installing git and Python 3")
    subprocess.check_call("brew install git || brew upgrade git", shell=True)
    subprocess.check_call("brew install python || brew upgrade python", shell=True)
    subprocess.check_call(['pip3', 'install', '--upgrade', 'click'])  # required for 'setup'

    setup_path = os.path.expanduser('~/setup')
    if os.path.exists(setup_path):
        print("Setup location already exists, updating")
        subprocess.check_call(['git', 'pull'], cwd=setup_path)
    else:
        print("Checking out setup repo")
        subprocess.check_call(['git', 'clone', REPO_URL], cwd=os.path.dirname(setup_path))

    setup_exe = os.path.join(setup_path, 'HOME/bin/setup')

    print("Installing all the things")
    subprocess.check_call([setup_exe, 'init'])
    print("Done installing all the things. Restart your terminal.")

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from __future__ import unicode_literals, division, absolute_import
import re
import logging
import os
from flexget.entry import Entry
from flexget.utils.cached_input import cached
from flexget.plugin import register_plugin, internet

log = logging.getLogger('regexp_parse')

class RegexpParse(object):
    """This plugin is designed to take input from a web resource or a file.
    It then parses the text via regexps supplied in the config file.

    source: is a file or url to get the data from. You can specify a username:password

    sections: Takes a list of dicts that contain regexps to split the data up into sections.
    The regexps listed here are used by find all so every matching string in the data will be
    a valid section.

    keys: hold the keys that will be set in the entries

    key:
      regexps: a list of dicts that hold regexps. The key is set to the first string that matches
      any of the regexps listed. The regexps are evaluated in the order they are supplied so if a
      string matches the first regexp none of the others in the list will be used.

      required: a boolean that when set to true will only allow entries that contain this key
      onto the next stage. url and title are always required no matter what you do (part of flexget)

      #TODO: consider adding a set field that will allow you to set the field if no regexps match

      #TODO: consider a mode field that allows a growing list for a field instead of just setting to
            # first match

    Example config

    regexp_parse:
      source: http://username:password@ezrss.it/feed/
      sections:
        - {regexp: "(?<=<item>).*?(?=</item>)", flags: "DOTALL"}

      keys:
        title:
          regexps:
            - {regexp: '(?<=<title><!\[CDATA\[).*?(?=\]\]></title>)'} #comment
        url:
          regexps:
            - {regexp: "magnet:.*?(?=])"}
        custom_field:
          regexps:
            - {regexp: "custom regexps", flags: "flags for python regexps see python regexps docs"}
          required: False
        custom_field2:
          regexps:
            - {regexp: 'first custom regexps'}
            - {regexp: 'can't find first regexp so try this one'}
    """

    #dict used to convert string values of regexp flags to int
    FLAG_VALUES = {'DEBUG': re.DEBUG,
                   'I': re.I,
                   'IGNORECASE': re.IGNORECASE,
                   'L': re.L,
                   'LOCALE': re.LOCALE,
                   'M': re.M,
                   'MULTILINE': re.MULTILINE,
                   'S': re.S,
                   'DOTALL': re.DOTALL,
                   'U': re.U,
                   'UNICODE': re.UNICODE,
                   'X': re.X,
                   'VERBOSE': re.VERBOSE
                   }

    def __init__(self):
        self.required = []

    def validator(self):
        from flexget import validator
        root = validator.factory('dict')

        root.accept('url', key='source', required=True)
        root.accept('file', key='source', required=True)

        #sections to divied source into
        sections_regexp_lists = root.accept('list', key='sections')
        section_regexp_list = sections_regexp_lists.accept('dict', required=True)
        section_regexp_list.accept('regexp', key='regexp', required=True)
        section_regexp_list.accept('text', key='flags')

        keys = root.accept('dict', key='keys', required=True)

        #required key need to specify for validator
        title = keys.accept('dict', key='title', required=True)
        title.accept('boolean', key='required')
        regexp_list = title.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        #required key need to specify for validator
        url = keys.accept_any_key('dict', key='url', required=True)
        url.accept('boolean', key='required')
        regexp_list = url.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        #accept any other key the user wants to use
        key = keys.accept_any_key('dict')
        key.accept('boolean', key='required')
        regexp_list = key.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        return root

    def flagstr_to_flags(self, flag_str):
        """turns a comma seperated list of flags into the int value."""
        COMBIND_FLAGS = 0
        split_flags = flag_str.split(',')
        for flag in split_flags:
            COMBIND_FLAGS = COMBIND_FLAGS | RegexpParse.FLAG_VALUES[flag.strip()]
        return COMBIND_FLAGS

    def compile_regexp_dict_list(self, re_list):
        """turns a list of dicts containing regexps information into a list of compiled regexps."""
        compiled_regexps = []
        for dic in re_list:
            flags = 0
            if 'flags' in dic:
                flags = self.flagstr_to_flags(dic['flags'])
            compiled_regexps.append(re.compile(dic['regexp'], flags))
        return compiled_regexps

    def isvalid(self, entry):
        """checks to make sure that all required fields are present in the entry."""
        for key in self.required:
            if key not in entry:
                return False
        return entry.isvalid()

    @cached('text')
    @internet(log)
    def on_task_input(self, task, config):

        entries = []

        url = config['source']

        #if it's a file open it and read into content
        if os.path.isfile(os.path.expanduser(url)):
            content = open(url).read()
        #else use requests to get the data
        else:
            content = task.requests.get(url).text

        sections = []
        seperators = config.get('sections')
        if seperators:
            for sep in seperators:
                flags = 0
                if 'flags' in sep:
                    flags = self.flagstr_to_flags(sep['flags'])
                sections.extend(re.findall(sep['regexp'], content, flags))

        #no seperators just do work on the whole content
        else:
            sections.append(content)

        #holds all the regex in a dict for the field they are trying to fill
        key_to_regexps = {}

        #put every key in keys into the rey_to_regexps list
        for key, value in config['keys'].iteritems():
            key_to_regexps[key] = self.compile_regexp_dict_list(value['regexps'])
            if 'required' in value and value['required']:
                self.required.append(key)

        entries = []
        for section in sections:
            entry = Entry()
            for key, regexps in key_to_regexps.iteritems():
                for regexp in regexps:
                    m = regexp.search(section)
                    if m:
                        entry[key] = m.group(0)
                        break
            if self.isvalid(entry):
                entries.append(entry)

        return entries

register_plugin(RegexpParse, 'regexp_parse', api_ver=2)
<commit_msg>Update regexp_parse.py<commit_after>from __future__ import unicode_literals, division, absolute_import
import re
import logging
import os
from flexget.entry import Entry
from flexget.utils.cached_input import cached
from flexget.plugin import register_plugin, internet

log = logging.getLogger('regexp_parse')

class RegexpParse(object):
    """This plugin is designed to take input from a web resource or a file.
    It then parses the text via regexps supplied in the config file.

    source: is a file or url to get the data from. You can specify a username:password

    sections: Takes a list of dicts that contain regexps to split the data up into sections.
    The regexps listed here are used by find all so every matching string in the data will be
    a valid section.

    keys: hold the keys that will be set in the entries

    key:
      regexps: a list of dicts that hold regexps. The key is set to the first string that matches
      any of the regexps listed. The regexps are evaluated in the order they are supplied so if a
      string matches the first regexp none of the others in the list will be used.

      required: a boolean that when set to true will only allow entries that contain this key
      onto the next stage. url and title are always required no matter what you do (part of flexget)

      #TODO: consider adding a set field that will allow you to set the field if no regexps match

      #TODO: consider a mode field that allows a growing list for a field instead of just setting to
            # first match

    Example config

    regexp_parse:
      source: http://username:password@ezrss.it/feed/
      sections:
        - {regexp: "(?<=<item>).*?(?=</item>)", flags: "DOTALL,IGNORECASE"}

      keys:
        title:
          regexps:
            - {regexp: '(?<=<title><!\[CDATA\[).*?(?=\]\]></title>)'} #comment
        url:
          regexps:
            - {regexp: "magnet:.*?(?=])"}
        custom_field:
          regexps:
            - {regexp: "custom regexps", flags: "comma seperated list of flags (see python regex docs)"}
          required: False
        custom_field2:
          regexps:
            - {regexp: 'first custom regexps'}
            - {regexp: 'can't find first regexp so try this one'}
    """

    #dict used to convert string values of regexp flags to int
    FLAG_VALUES = {'DEBUG': re.DEBUG,
                   'I': re.I,
                   'IGNORECASE': re.IGNORECASE,
                   'L': re.L,
                   'LOCALE': re.LOCALE,
                   'M': re.M,
                   'MULTILINE': re.MULTILINE,
                   'S': re.S,
                   'DOTALL': re.DOTALL,
                   'U': re.U,
                   'UNICODE': re.UNICODE,
                   'X': re.X,
                   'VERBOSE': re.VERBOSE
                   }

    def __init__(self):
        self.required = []

    def validator(self):
        from flexget import validator
        root = validator.factory('dict')

        root.accept('url', key='source', required=True)
        root.accept('file', key='source', required=True)

        #sections to divied source into
        sections_regexp_lists = root.accept('list', key='sections')
        section_regexp_list = sections_regexp_lists.accept('dict', required=True)
        section_regexp_list.accept('regexp', key='regexp', required=True)
        section_regexp_list.accept('text', key='flags')

        keys = root.accept('dict', key='keys', required=True)

        #required key need to specify for validator
        title = keys.accept('dict', key='title', required=True)
        title.accept('boolean', key='required')
        regexp_list = title.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        #required key need to specify for validator
        url = keys.accept_any_key('dict', key='url', required=True)
        url.accept('boolean', key='required')
        regexp_list = url.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        #accept any other key the user wants to use
        key = keys.accept_any_key('dict')
        key.accept('boolean', key='required')
        regexp_list = key.accept('list', key='regexps', required=True)
        regexp = regexp_list.accept('dict', required=True)
        regexp.accept('regexp', key='regexp', required=True)
        regexp.accept('text', key='flags')

        return root

    def flagstr_to_flags(self, flag_str):
        """turns a comma seperated list of flags into the int value."""
        COMBIND_FLAGS = 0
        split_flags = flag_str.split(',')
        for flag in split_flags:
            COMBIND_FLAGS = COMBIND_FLAGS | RegexpParse.FLAG_VALUES[flag.strip()]
        return COMBIND_FLAGS

    def compile_regexp_dict_list(self, re_list):
        """turns a list of dicts containing regexps information into a list of compiled regexps."""
        compiled_regexps = []
        for dic in re_list:
            flags = 0
            if 'flags' in dic:
                flags = self.flagstr_to_flags(dic['flags'])
            compiled_regexps.append(re.compile(dic['regexp'], flags))
        return compiled_regexps

    def isvalid(self, entry):
        """checks to make sure that all required fields are present in the entry."""
        for key in self.required:
            if key not in entry:
                return False
        return entry.isvalid()

    @cached('text')
    @internet(log)
    def on_task_input(self, task, config):

        entries = []

        url = config['source']

        #if it's a file open it and read into content
        if os.path.isfile(os.path.expanduser(url)):
            content = open(url).read()
        #else use requests to get the data
        else:
            content = task.requests.get(url).text

        sections = []
        seperators = config.get('sections')
        if seperators:
            for sep in seperators:
                flags = 0
                if 'flags' in sep:
                    flags = self.flagstr_to_flags(sep['flags'])
                sections.extend(re.findall(sep['regexp'], content, flags))

        #no seperators just do work on the whole content
        else:
            sections.append(content)

        #holds all the regex in a dict for the field they are trying to fill
        key_to_regexps = {}

        #put every key in keys into the rey_to_regexps list
        for key, value in config['keys'].iteritems():
            key_to_regexps[key] = self.compile_regexp_dict_list(value['regexps'])
            if 'required' in value and value['required']:
                self.required.append(key)

        entries = []
        for section in sections:
            entry = Entry()
            for key, regexps in key_to_regexps.iteritems():
                for regexp in regexps:
                    m = regexp.search(section)
                    if m:
                        entry[key] = m.group(0)
                        break
            if self.isvalid(entry):
                entries.append(entry)

        return entries

register_plugin(RegexpParse, 'regexp_parse', api_ver=2)
<|endoftext|>


<commit_before>"""
Functions for plotting datasets nicely.
"""
from itertools import cycle
from collections import OrderedDict
from math import log
from itertools import repeat
import numpy as np

# -------------------------------------------------------------------------- #
# Plots with matplotlib                                                      #
# -------------------------------------------------------------------------- #

def mpl_markers():
    marker_dict = OrderedDict([
        ("o", "circle"),
        ("x", "x"),
        ("D", "diamond"),
        ("+", "plus"),
        ("s", "square"),
        (".", "point"),
        ("^", "triangle_up"),
        ("3", "tri_left"),
        (">", "triangle_right"),
        ("d", "thin_diamond"),
        ("*", "star"),
        ("v", "triangle_down"),
        ("|", "vline"),
        ("1", "tri_down"),
        ("p", "pentagon"),
        (",", "pixel"),
        ("2", "tri_up"),
        ("<", "triangle_left"),
        ("h", "hexagon1"),
        ("4", "tri_right"),
        (0, "tickleft"),
        (2, "tickup"),
        (3, "tickdown"),
        (4, "caretleft"),
        ("_", "hline"),
        (5, "caretright"),
        ("H", "hexagon2"),
        (1, "tickright"),
        (6, "caretup"),
        ("8", "octagon"),
        (7, "caretdown"),
        ])
    marker_keys = [*marker_dict.keys()]
    return marker_keys

def mplot(x, y_i, fignum=1, logx=False, logy=False,
          xlims=None, ylims=None, markers=True,
          color=False, colormap="viridis", **kwargs):
    """ Function for automatically plotting multiple sets of data
    using matplot lib. """
    import matplotlib.pyplot as plt
    y_i = np.array(np.squeeze(y_i), ndmin=2)
    if np.size(x) == y_i.shape[0]:
        y_i = np.transpose(y_i)
    n_y = y_i.shape[0]
    fig = plt.figure(fignum, figsize=(8, 6), dpi=100)
    axes = fig.add_axes([0.1, 0.1, 0.85, 0.8])

    if markers:
        mrkrs = cycle(mpl_markers())
    else:
        repeat(None)

    if color:
        from matplotlib import cm
        cmap = getattr(cm, colormap)
        cns = np.linspace(0, 1, n_y)
        cols = [cmap(cn, 1) for cn in cns]
    else:
        cols = repeat(None)

    for y, col, mrkr in zip(y_i, cols, mrkrs):
        axes.plot(x, y, ".-", c=col, marker=mrkr, **kwargs)
    xlims = (np.min(x), np.max(x)) if xlims is None else xlims
    ylims = (np.min(y_i), np.max(y_i)) if ylims is None else ylims
    axes.set_xlim(xlims)
    axes.set_ylim(ylims)
    axes.set_xscale("log" if logx else "linear")
    axes.set_yscale("log" if logy else "linear")
    return fig

# -------------------------------------------------------------------------- #
# Plots with matplotlib and xarray                                           #
# -------------------------------------------------------------------------- #

def xmlineplot(ds, y_coo, x_coo, z_coo, title=None, legend=None,
               xlabel=None, ylabel=None, zlabel=None,
               xlims=None, ylims=None, markers=True,
               vlines=None, hlines=None, color=False, colormap="viridis",
               fignum=1, logx=False, logy=False, **kwargs):
    """
    Function for automatically plotting multiple sets of data
    using matplotlib and xarray.
    """
    import matplotlib.pyplot as plt
    fig = plt.figure(fignum, figsize=(8, 6), dpi=100)
    axes = fig.add_axes([0.1, 0.1, 0.85, 0.8],
                        title=("" if title is None else title))

    if markers:
        mrkrs = cycle(mpl_markers())
    else:
        repeat(None)

    if color:
        from matplotlib import cm
        cmap = getattr(cm, colormap)
        zmin = ds[z_coo].values.min()
        zmax = ds[z_coo].values.max()
        cols = [cmap(1 - (z-zmin)/(zmax-zmin))
                for z in ds[z_coo].values]
    else:
        cols = repeat(None)

    for z, col, mrkr in zip(ds[z_coo].data, cols, mrkrs):
        x = ds.loc[{z_coo: z}][x_coo].data.flatten()
        y = ds.loc[{z_coo: z}][y_coo].data.flatten()
        axes.plot(x, y, ".-", c=col, lw=1.3, marker=mrkr,
                  label=str(z), zorder=3, **kwargs)
    axes.set_xscale("log" if logx else "linear")
    axes.set_yscale("log" if logy else "linear")
    if vlines is not None:
        for x in vlines:
            axes.axvline(x)
    if hlines is not None:
        for y in hlines:
            axes.axhline(y)
    axes.grid(True, color="0.666")
    xlims = (ds[x_coo].min(), ds[x_coo].max()) if xlims is None else xlims
    ylims = (ds[y_coo].min(), ds[y_coo].max()) if ylims is None else ylims
    axes.set_xlim(xlims)
    axes.set_ylim(ylims)
    axes.set_xlabel(x_coo if xlabel is None else xlabel)
    axes.set_ylabel(y_coo if ylabel is None else ylabel)
    if legend or not (legend is False or len(ds[z_coo]) > 20):
        axes.legend(title=(z_coo if zlabel is None else zlabel), loc="best")
    return fig

# -------------------------------------------------------------------------- #
# Plots with plotly                                                          #
# -------------------------------------------------------------------------- #
def ishow(figs, nb=True, **kwargs):
    """ Show multiple plotly figures in notebook or on web. """
    if isinstance(figs, (list, tuple)):
        fig_main = figs[0]
        for fig in figs[1:]:
            fig_main['data'] += fig['data']
    else:
        fig_main = figs
    if nb:
        from plotly.offline import init_notebook_mode
        from plotly.offline import iplot as plot
        init_notebook_mode()
    else:
        from plotly.plotly import plot
    plot(fig_main, **kwargs)

def iplot(x, y_i, logx=False, logy=False, color=False, colormap="viridis",
          return_fig=False, nb=True, go_dict={}, ly_dict={},
          **kwargs):
    """ Multi line plot with plotly. """
    from plotly.graph_objs import Scatter
    y_i = np.array(np.squeeze(y_i), ndmin=2)
    if np.size(x) == y_i.shape[0]:
        y_i = np.transpose(y_i)
    n_y = y_i.shape[0]
    if color:
        import matplotlib.cm as cm
        cmap = getattr(cm, colormap)
        cns = np.linspace(0, 1, n_y)
        cols = ["rgba" + str(cmap(cn)) for cn in cns]
    else:
        cols = repeat(None)
    traces = [Scatter({"x": x, "y": y, "line": {"color": col},
                       "marker": {"color": col},  **go_dict})
              for y, col in zip(y_i, cols)]
    layout = {"width": 750, "height": 600, "showlegend": False,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logy else "linear"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def iscatter(x, y, cols=None, logx=False, logy=False, nb=True,
             return_fig=False, ly_dict={}, **kwargs):
    from plotly.graph_objs import Scatter, Marker
    mkr = Marker({"color": cols, "opacity": 0.9,
                  "colorscale": "Portland", "showscale": cols is not None})
    traces = [Scatter({"x": x, "y": y, "mode": "markers", "marker": mkr})]
    layout = {"width": 700, "height": 700, "showlegend": False,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logy else "linear"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def ihist(xs, nb=True, go_dict={}, ly_dict={}, return_fig=False,
          **kwargs):
    """ Histogram plot with plotly. """
    from plotly.graph_objs import Histogram
    traces = [Histogram({"x": x, **go_dict}) for x in xs]
    layout = {"width": 750, "height": 600,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

# -------------------------------------------------------------------------- #
# Plots with plotly and xarray                                               #
# -------------------------------------------------------------------------- #

def iheatmap(ds, data_name, x_coo, y_coo, colormap="Portland",
             go_dict={}, ly_dict={}, nb=True, return_fig=False,
             **kwargs):
    """
    Automatic 2D-Heatmap plot using plotly.
    """
    from plotly.graph_objs import Heatmap
    traces = Heatmap({"z": (ds[data_name]
                            .dropna(x_coo, how="all")
                            .dropna(y_coo, how="all")
                            .squeeze()
                            .transpose(y_coo, x_coo)
                            .data),
                      "x": ds.coords[x_coo].values,
                      "y": ds.coords[y_coo].values,
                      "colorscale": colormap,
                      "colorbar": {"title": data_name}, **go_dict})
    layout = {"height": 600, "width": 650,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "outside", "title": x_coo},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "outside", "title": y_coo}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def ilineplot(ds, data_name, x_coo, z_coo=None, logx=False, logy=False,
              erry=None, errx=None, nb=True, color=False, colormap="viridis",
              log_color=False, legend=None, traces=[], go_dict={},
              ly_dict={}, return_fig=False, **kwargs):
    # TODO: add hlines, vlines, xlims, ylims, title
    from plotly.graph_objs import Scatter
    if z_coo is None:
        traces = [Scatter({
                    "x": ds[x_coo].values,
                    "y": ds[data_name].values.flatten(), **go_dict})]
    else:
        cols = (calc_colors(ds, z_coo, colormap=colormap, log_scale=log_color)
                if color else repeat(None))
        traces = [Scatter({
                    "x": ds.loc[{z_coo: z}][x_coo].values.flatten(),
                    "y": ds.loc[{z_coo: z}][data_name].values.flatten(),
                    "name": str(z), "line": {"color": col},
                    "marker": {"color": col}, **go_dict})
                  for z, col in zip(ds[z_coo].values, cols)]
    layout = {"width": 750, "height": 600,
              "xaxis": {"showline": True, "title": x_coo,
                        "mirror": "ticks", "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "title": data_name,
                        "mirror": "ticks", "ticks": "inside",
                        "type": "log" if logy else "linear"},
              "showlegend": legend or not (legend is False or z_coo is None or
                                           len(ds[z_coo]) > 20), **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def calc_colors(ds, z_coo, colormap="viridis", log_scale=False):
    import matplotlib.cm as cm
    cmap = getattr(cm, colormap)
    zmin = ds[z_coo].values.min()
    zmax = ds[z_coo].values.max()
    if log_scale:
        cols = ["rgba" + str(cmap(1 -
                                  (log(z)-log(zmin))/(log(zmax)-log(zmin))))
                for z in ds[z_coo].values]
    else:
        cols = ["rgba" + str(cmap(1 - (z - zmin) / (zmax - zmin)))
                for z in ds[z_coo].values]
    return cols
<commit_msg>style changes to xmlineplot<commit_after>"""
Functions for plotting datasets nicely.
"""
from itertools import cycle
from collections import OrderedDict
from math import log
from itertools import repeat
import numpy as np

# -------------------------------------------------------------------------- #
# Plots with matplotlib                                                      #
# -------------------------------------------------------------------------- #

def mpl_markers():
    marker_dict = OrderedDict([
        ("o", "circle"),
        ("x", "x"),
        ("D", "diamond"),
        ("+", "plus"),
        ("s", "square"),
        (".", "point"),
        ("^", "triangle_up"),
        ("3", "tri_left"),
        (">", "triangle_right"),
        ("d", "thin_diamond"),
        ("*", "star"),
        ("v", "triangle_down"),
        ("|", "vline"),
        ("1", "tri_down"),
        ("p", "pentagon"),
        (",", "pixel"),
        ("2", "tri_up"),
        ("<", "triangle_left"),
        ("h", "hexagon1"),
        ("4", "tri_right"),
        (0, "tickleft"),
        (2, "tickup"),
        (3, "tickdown"),
        (4, "caretleft"),
        ("_", "hline"),
        (5, "caretright"),
        ("H", "hexagon2"),
        (1, "tickright"),
        (6, "caretup"),
        ("8", "octagon"),
        (7, "caretdown"),
        ])
    marker_keys = [*marker_dict.keys()]
    return marker_keys

def mplot(x, y_i, fignum=1, logx=False, logy=False,
          xlims=None, ylims=None, markers=True,
          color=False, colormap="viridis", **kwargs):
    """ Function for automatically plotting multiple sets of data
    using matplot lib. """
    import matplotlib.pyplot as plt
    y_i = np.array(np.squeeze(y_i), ndmin=2)
    if np.size(x) == y_i.shape[0]:
        y_i = np.transpose(y_i)
    n_y = y_i.shape[0]
    fig = plt.figure(fignum, figsize=(8, 6), dpi=100)
    axes = fig.add_axes([0.1, 0.1, 0.85, 0.8])

    if markers:
        mrkrs = cycle(mpl_markers())
    else:
        repeat(None)

    if color:
        from matplotlib import cm
        cmap = getattr(cm, colormap)
        cns = np.linspace(0, 1, n_y)
        cols = [cmap(cn, 1) for cn in cns]
    else:
        cols = repeat(None)

    for y, col, mrkr in zip(y_i, cols, mrkrs):
        axes.plot(x, y, ".-", c=col, marker=mrkr, **kwargs)
    xlims = (np.min(x), np.max(x)) if xlims is None else xlims
    ylims = (np.min(y_i), np.max(y_i)) if ylims is None else ylims
    axes.set_xlim(xlims)
    axes.set_ylim(ylims)
    axes.set_xscale("log" if logx else "linear")
    axes.set_yscale("log" if logy else "linear")
    return fig

# -------------------------------------------------------------------------- #
# Plots with matplotlib and xarray                                           #
# -------------------------------------------------------------------------- #

def xmlineplot(ds, y_coo, x_coo, z_coo, title=None, legend=None,
               xlabel=None, ylabel=None, zlabel=None,
               xlims=None, ylims=None, markers=True, padding=0.0,
               vlines=None, hlines=None, color=False, colormap="viridis",
               fignum=1, logx=False, logy=False, **kwargs):
    """
    Function for automatically plotting multiple sets of data
    using matplotlib and xarray.
    """
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    mpl.rc('font', family='CMU Serif')
    fig = plt.figure(fignum, figsize=(8, 6), dpi=100)
    axes = fig.add_axes([0.15, 0.15, 0.8, 0.75],
                        title=("" if title is None else title))
    axes.tick_params(labelsize=16)

    if markers:
        mrkrs = cycle(mpl_markers())
    else:
        repeat(None)

    if color:
        from matplotlib import cm
        cmap = getattr(cm, colormap)
        zmin, zmax = ds[z_coo].values.min(), ds[z_coo].values.max()
        cols = [cmap(1 - (z-zmin)/(zmax-zmin)) for z in ds[z_coo].values]
    else:
        cols = repeat(None)

    # Add data plots
    for z, col, mrkr in zip(ds[z_coo].data, cols, mrkrs):
        x = ds.loc[{z_coo: z}][x_coo].data.flatten()
        y = ds.loc[{z_coo: z}][y_coo].data.flatten()
        axes.plot(x, y, ".-", c=col, lw=1.3, marker=mrkr,
                  label=str(z), zorder=3, **kwargs)

    if vlines is not None:
        for x in vlines:
            axes.axvline(x)
    if hlines is not None:
        for y in hlines:
            axes.axhline(y)

    axes.grid(True, color="0.666")

    if xlims is None:
        xmax = ds[x_coo].max()
        xmin = ds[x_coo].min()
        xrange = xmax - xmin
        xlims = (xmin - padding * xrange, xmax + padding * xrange)
    axes.set_xlim(xlims)
    if ylims is None:
        ymax = ds[y_coo].max()
        ymin = ds[y_coo].min()
        yrange = ymax - ymin
        ylims = (ymin - padding * yrange, ymax + padding * yrange)
    axes.set_ylim(ylims)

    axes.set_xscale("log" if logx else "linear")
    axes.set_yscale("log" if logy else "linear")
    axes.set_xlabel(x_coo if xlabel is None else xlabel, fontsize=20)
    axes.set_ylabel(y_coo if ylabel is None else ylabel, fontsize=20)

    if legend or not (legend is False or len(ds[z_coo]) > 20):
        legend = axes.legend(title=(z_coo if zlabel is None else zlabel),
                             loc="best", fontsize=16, frameon=False)
        legend.get_title().set_fontsize(20)
    return fig

# -------------------------------------------------------------------------- #
# Plots with plotly                                                          #
# -------------------------------------------------------------------------- #
def ishow(figs, nb=True, **kwargs):
    """ Show multiple plotly figures in notebook or on web. """
    if isinstance(figs, (list, tuple)):
        fig_main = figs[0]
        for fig in figs[1:]:
            fig_main['data'] += fig['data']
    else:
        fig_main = figs
    if nb:
        from plotly.offline import init_notebook_mode
        from plotly.offline import iplot as plot
        init_notebook_mode()
    else:
        from plotly.plotly import plot
    plot(fig_main, **kwargs)

def iplot(x, y_i, logx=False, logy=False, color=False, colormap="viridis",
          return_fig=False, nb=True, go_dict={}, ly_dict={},
          **kwargs):
    """ Multi line plot with plotly. """
    from plotly.graph_objs import Scatter
    y_i = np.array(np.squeeze(y_i), ndmin=2)
    if np.size(x) == y_i.shape[0]:
        y_i = np.transpose(y_i)
    n_y = y_i.shape[0]
    if color:
        import matplotlib.cm as cm
        cmap = getattr(cm, colormap)
        cns = np.linspace(0, 1, n_y)
        cols = ["rgba" + str(cmap(cn)) for cn in cns]
    else:
        cols = repeat(None)
    traces = [Scatter({"x": x, "y": y, "line": {"color": col},
                       "marker": {"color": col},  **go_dict})
              for y, col in zip(y_i, cols)]
    layout = {"width": 750, "height": 600, "showlegend": False,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logy else "linear"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def iscatter(x, y, cols=None, logx=False, logy=False, nb=True,
             return_fig=False, ly_dict={}, **kwargs):
    from plotly.graph_objs import Scatter, Marker
    mkr = Marker({"color": cols, "opacity": 0.9,
                  "colorscale": "Portland", "showscale": cols is not None})
    traces = [Scatter({"x": x, "y": y, "mode": "markers", "marker": mkr})]
    layout = {"width": 700, "height": 700, "showlegend": False,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside",
                        "type": "log" if logy else "linear"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def ihist(xs, nb=True, go_dict={}, ly_dict={}, return_fig=False,
          **kwargs):
    """ Histogram plot with plotly. """
    from plotly.graph_objs import Histogram
    traces = [Histogram({"x": x, **go_dict}) for x in xs]
    layout = {"width": 750, "height": 600,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside"},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "inside"}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

# -------------------------------------------------------------------------- #
# Plots with plotly and xarray                                               #
# -------------------------------------------------------------------------- #

def iheatmap(ds, data_name, x_coo, y_coo, colormap="Portland",
             go_dict={}, ly_dict={}, nb=True, return_fig=False,
             **kwargs):
    """
    Automatic 2D-Heatmap plot using plotly.
    """
    from plotly.graph_objs import Heatmap
    traces = Heatmap({"z": (ds[data_name]
                            .dropna(x_coo, how="all")
                            .dropna(y_coo, how="all")
                            .squeeze()
                            .transpose(y_coo, x_coo)
                            .data),
                      "x": ds.coords[x_coo].values,
                      "y": ds.coords[y_coo].values,
                      "colorscale": colormap,
                      "colorbar": {"title": data_name}, **go_dict})
    layout = {"height": 600, "width": 650,
              "xaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "outside", "title": x_coo},
              "yaxis": {"showline": True, "mirror": "ticks",
                        "ticks": "outside", "title": y_coo}, **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def ilineplot(ds, data_name, x_coo, z_coo=None, logx=False, logy=False,
              erry=None, errx=None, nb=True, color=False, colormap="viridis",
              log_color=False, legend=None, traces=[], go_dict={},
              ly_dict={}, return_fig=False, **kwargs):
    # TODO: add hlines, vlines, xlims, ylims, title
    from plotly.graph_objs import Scatter
    if z_coo is None:
        traces = [Scatter({
                    "x": ds[x_coo].values,
                    "y": ds[data_name].values.flatten(), **go_dict})]
    else:
        cols = (calc_colors(ds, z_coo, colormap=colormap, log_scale=log_color)
                if color else repeat(None))
        traces = [Scatter({
                    "x": ds.loc[{z_coo: z}][x_coo].values.flatten(),
                    "y": ds.loc[{z_coo: z}][data_name].values.flatten(),
                    "name": str(z), "line": {"color": col},
                    "marker": {"color": col}, **go_dict})
                  for z, col in zip(ds[z_coo].values, cols)]
    layout = {"width": 750, "height": 600,
              "xaxis": {"showline": True, "title": x_coo,
                        "mirror": "ticks", "ticks": "inside",
                        "type": "log" if logx else "linear"},
              "yaxis": {"showline": True, "title": data_name,
                        "mirror": "ticks", "ticks": "inside",
                        "type": "log" if logy else "linear"},
              "showlegend": legend or not (legend is False or z_coo is None or
                                           len(ds[z_coo]) > 20), **ly_dict}
    fig = {"data": traces, "layout": layout}
    if return_fig:
        return fig
    ishow(fig, nb=nb, **kwargs)

def calc_colors(ds, z_coo, colormap="viridis", log_scale=False):
    import matplotlib.cm as cm
    cmap = getattr(cm, colormap)
    zmin = ds[z_coo].values.min()
    zmax = ds[z_coo].values.max()
    if log_scale:
        cols = ["rgba" + str(cmap(1 -
                                  (log(z)-log(zmin))/(log(zmax)-log(zmin))))
                for z in ds[z_coo].values]
    else:
        cols = ["rgba" + str(cmap(1 - (z - zmin) / (zmax - zmin)))
                for z in ds[z_coo].values]
    return cols
<|endoftext|>