<commit_before>from os.path import join, abspath, dirname, getsize, curdir
from netCDF4 import Dataset
from flyingpigeon import config
import logging
LOGGER = logging.getLogger("PYWPS")

def has_Lambert_Conformal(resource):
    """
    Check if grid is organised as Lambert_Conformal

    :param resource: file to be checked

    :return Boolean: True/False
    """
    if type(resource) != list:
        resource = [resource]
    for nc in resource:
        ds = Dataset(nc)
        if 'Lambert_Conformal' not in ds.variables.keys():
            return False
    return True

def call(resource=[], variable=None, dimension_map=None, calc=None,
         calc_grouping=None, conform_units_to=None, crs=None, memory_limit=None,  prefix=None,
         regrid_destination=None, regrid_options='bil', level_range=None,
         geom=None, output_format_options=None, search_radius_mult=2.,
         select_nearest=False, select_ugid=None, spatial_wrapping=None,
         t_calendar=None, time_region=None,
         time_range=None, dir_output=None, output_format='nc'):
    '''
    ocgis operation call

    :param resource:
    :param variable: variable in the input file to be picked
    :param dimension_map: dimension map in case of unconventional storage of data
    :param calc: ocgis calc syntax for calculation partion
    :param calc_grouping: time aggregate grouping
    :param conform_units_to:
    :param crs: coordinate reference system
    :param memory_limit: limit the amount of data to be loaded into the memory at once \
        if None (default) free memory is detected by birdhouse
    :param level_range: subset of given levels
    :param prefix: string for the file base name
    :param regrid_destination: file path with netCDF file with grid for output file
    :param geom: name of shapefile stored in birdhouse shape cabinet
    :param output_format_options: output options for netCDF e.g compression level()
    :param regrid_destination: file containing the targed grid (griddes.txt or netCDF file)
    :param regrid_options: methods for regridding:
                          'bil' = Bilinear interpolation
                          'bic' = Bicubic interpolation
                          'dis' = Distance-weighted average remapping
                          'nn' = nearest neighbour
                          'con' = First-order conservative remapping
                          'laf' = largest area fraction reamapping
    :param search_radius_mult: search radius for point geometries. All included gridboxes will be returned
    :param select_nearest: nearest neighbour selection for point geometries
    :param select_ugid: ugid for appropriate polygons
    :param spatial_wrapping: how to handle coordinates in case of subsets, options: None (default), 'wrap', 'unwrap'
    :param time_region: select single month
    :param time_range: sequence of two datetime.datetime objects to mark start and end point
    :param dir_output (default= curdir):
    :param output_format:
    :return: output file path
    '''
    LOGGER.info('Start ocgis module call function')
    from ocgis import OcgOperations, RequestDataset, env, DimensionMap, crs
    from ocgis.util.large_array import compute
    from datetime import datetime as dt
    from datetime import date as dd
    from datetime import time as dt_time
    import uuid

    # prepare the environment
    env.DIR_SHPCABINET = config.shapefiles_path()
    env.OVERWRITE = True
    # env.DIR_OUTPUT = dir_output
    # LOGGER.debug(' **** env.DIR_OUTPUT  = %s ' % env.DIR_OUTPUT)

    if dir_output is None:
        dir_output = abspath(curdir)

    # check time_range format:

    if time_range is not None:
        try:
            LOGGER.debug('time_range type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            LOGGER.debug('time_range= %s , %s ' % (time_range[0], time_range[1]))
            #if type(time_range[0] is 'datetime.date'):
            if (isinstance(time_range[0],dd) and not isinstance(time_range[0],dt)):
                time_range = [dt.combine(time_range[0], dt.min.time()),
                              dt.combine(time_range[1], dt.min.time())]
                # time_range = [dt.combine(time_range[0], dt_time(12,0)),
                #               dt.combine(time_range[1], dt_time(12,0))]
            LOGGER.debug('time_range changed to type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            LOGGER.debug('time_range changed to= %s , %s ' % (time_range[0], time_range[1]))
        except:
            LOGGER.exception('failed to confert data to datetime')

    if spatial_wrapping == 'wrap':
        spatial_reorder = True
    else:
        spatial_reorder = False
    LOGGER.debug('spatial_reorder: %s and spatial_wrapping: %s ' % (spatial_reorder, spatial_wrapping))

    if prefix is None:
        prefix = str(uuid.uuid1())
        env.PREFIX = prefix
    #
    # if output_format_options is False:
    #     output_format_options = None
    # elif output_format_options is True:
    #     output_format_options = {'data_model': 'NETCDF4',  # NETCDF4_CLASSIC
    #                              'variable_kwargs': {'zlib': True, 'complevel': 9}}
    # else:
    if output_format_options is not None:
        LOGGER.info('output_format_options are set to %s ' % (output_format_options))

    if type(resource) != list:
        resource = list([resource])
    # execute ocgis
    LOGGER.info('Execute ocgis module call function')

    # # needed for some AFR-44 data
    # dimension_map = DimensionMap()
    # dimension_map.set_variable('x', 'lon', dimension='rlon')
    # dimension_map.set_variable('y', 'lat', dimension='rlat')
    # dimension_map.set_variable('time', 'time', dimension='time')
    # crs=crs.Spherical()

    try:
        LOGGER.debug('call module curdir = %s ' % abspath(curdir))
        rd = RequestDataset(resource, variable=variable, level_range=level_range,
                            dimension_map=dimension_map, conform_units_to=conform_units_to,
                            time_region=time_region, t_calendar=t_calendar, time_range=time_range) # , crs=crs)

        from ocgis.constants import DimensionMapKey
        rd.dimension_map.set_bounds(DimensionMapKey.TIME, None)

        ops = OcgOperations(dataset=rd,
                            output_format_options=output_format_options,
                            dir_output=dir_output,
                            spatial_wrapping=spatial_wrapping,
                            spatial_reorder=spatial_reorder,
                            # regrid_destination=rd_regrid,
                            # options=options,
                            calc=calc,
                            calc_grouping=calc_grouping,
                            geom=geom,
                            output_format=output_format,
                            prefix=prefix,
                            search_radius_mult=search_radius_mult,
                            select_nearest=select_nearest,
                            select_ugid=select_ugid,
                            add_auxiliary_files=False)
        LOGGER.info('OcgOperations set')
    except:
        LOGGER.exception('failed to setup OcgOperations')
        return None

    try:
        LOGGER.info('ocgis module call as ops.execute()')
        geom_file = ops.execute()
    except:
        LOGGER.exception('failed to execute ocgis operation')
        return None

    #
    # try:
    #     from numpy import sqrt
    #     from flyingpigeon.utils import FreeMemory
    #
    #     if memory_limit is None:
    #         f = FreeMemory()
    #         mem_kb = f.user_free
    #         mem_mb = mem_kb / 1024.
    #         mem_limit = mem_mb / 2.  # set limit to half of the free memory
    #     else:
    #         mem_limit = memory_limit
    #
    #     if mem_limit >= 1024. * 4:
    #         mem_limit = 1024. * 4
    #         # 475.0 MB for openDAP
    #
    #     LOGGER.info('memory_limit = %s Mb' % (mem_limit))
    #
    #     data_kb = ops.get_base_request_size()['total']
    #     data_mb = data_kb / 1024.
    #
    #     # data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
    #     LOGGER.info('data_mb  = %s Mb' % (data_mb))
    #
    #     if data_mb <= mem_limit:  # input is smaler than the half of free memory size
    #         try:
    #             LOGGER.info('ocgis module call as ops.execute()')
    #             geom_file = ops.execute()
    #         except Exception as e:
    #             LOGGER.debug('failed to execute ocgis operation')
    #             raise
    #             return None
    #
    #     else:
    #         ##########################
    #         # calcultion of chunk size
    #         ##########################
    #         try:
    #             size = ops.get_base_request_size()
    #             nb_time_coordinates_rd = size['variables'][variable]['temporal']['shape'][0]
    #             element_in_kb = size['total']/reduce(lambda x, y: x*y, size['variables'][variable]['value']['shape'])
    #             element_in_mb = element_in_kb / 1024.
    #             tile_dim = sqrt(mem_limit/(element_in_mb*nb_time_coordinates_rd))  # maximum chunk size
    #
    #             LOGGER.info('ocgis module call compute with chunks')
    #             if calc is None:
    #                 calc = '%s=%s*1' % (variable, variable)
    #                 LOGGER.info('calc set to = %s ' % calc)
    #             ops = OcgOperations(dataset=rd,
    #                                 output_format_options=output_format_options,
    #                                 dir_output=dir_output,
    #                                 spatial_wrapping=spatial_wrapping,
    #                                 spatial_reorder=spatial_reorder,
    #                                 # regrid_destination=rd_regrid,
    #                                 # options=options,
    #                                 calc=calc,
    #                                 calc_grouping=calc_grouping,
    #                                 geom=geom,
    #                                 output_format=output_format,
    #                                 prefix=prefix,
    #                                 search_radius_mult=search_radius_mult,
    #                                 select_nearest=select_nearest,
    #                                 select_ugid=select_ugid,
    #                                 add_auxiliary_files=False)
    #             geom_file = compute(ops, tile_dimension=int(tile_dim), verbose=True)
    #             print 'ocgis calculated'
    #         except Exception as e:
    #             LOGGER.debug('failed to compute ocgis with chunks')
    #             raise
    #             return None
    #     LOGGER.info('Succeeded with ocgis module call function')
    # except:
    #     LOGGER.exception('failed to compare dataload with free memory, calling as execute instead')

    ############################################
    # remapping according to regrid informations
    ############################################
    if regrid_destination is not None:
        try:
            from tempfile import mkstemp
            from cdo import Cdo
            cdo = Cdo()
            output = '%s.nc' % uuid.uuid1()
            remap = 'remap%s' % regrid_options
            call = [op for op in dir(cdo) if remap in op]
            cmd = "output = cdo.%s('%s',input='%s', output='%s')" \
                  % (str(call[0]), regrid_destination, geom_file, output)
            exec cmd
        except Exception as e:
            LOGGER.debug('failed to remap')
            raise
            return None
    else:
        output = geom_file

    # try:
    #     from flyingpigeon.utils import unrotate_pole
    #     lat, lon = unrotate_pole(output)
    # except:
    #     LOGGER.exception('failed to unrotate pole')
    return output

def eval_timerange(resource, time_range):
    """
    quality checker if given time_range is covered by timesteps in resource files

    :param resource: input netCDF files
    :param time_range: start and end date of time range [datetime,datetime]

    :returns [datetime,datetime]: time_range
    """
    from flyingpigeon.utils import get_time

    LOGGER.info('time_range: %s' % time_range)

    if type(resource) != str:
        resource.sort()
    time = get_time(resource)
    start = time[0]
    end = time[-1]

    if (time_range[0] > start or time_range[0] < end):
        LOGGER.debug('time range start %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[0] = start
    LOGGER.debug('time_range start changed to first timestep of dataset')
    if (time_range[1] > end or time_range[1] < start):
        LOGGER.debug('time range end %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[1] = end
    LOGGER.debug('time_range end changed to last timestep of dataset')
    if (time_range[0] > time_range[1]):
        time_range = reversed(time_range)
        LOGGER.debug('time range reversed! start was later than end ')
    LOGGER.info('time range start and end set')
    return time_range

# # check memory load
# from os import stat
#   if memory_limit == None:
#     f = FreeMemory()
#     mem_kb = f.user_free
#     mem_mb = mem_kb / 1024.
#     mem_limit = mem_mb / 2. # set limit to half of the free memory
#   else:
#     mem_limit = memory_limit
#
#   if mem_limit >= 1024. * 4:
#     mem_limit = 1024. * 4
#     # 475.0 MB for openDAP
#
#   #if type(resource) == list :
#     #data_kb =  stat(resource[0]).st_size * len(resource)
#   #else:
#     #data_kb =  stat(resource).st_size
#   size = ops.get_base_request_size()['total']
#   data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
#   data_mb = data_kb / 1024.
#
#   if variable == None:
#     variable = rd.variable
#     LOGGER.info('%s as variable dedected' % (variable))
<commit_msg>add option for regriding<commit_after>from os.path import join, abspath, dirname, getsize, curdir
from netCDF4 import Dataset
from flyingpigeon import config
import logging
LOGGER = logging.getLogger("PYWPS")

def has_Lambert_Conformal(resource):
    """
    Check if grid is organised as Lambert_Conformal

    :param resource: file to be checked

    :return Boolean: True/False
    """
    if type(resource) != list:
        resource = [resource]
    for nc in resource:
        ds = Dataset(nc)
        if 'Lambert_Conformal' not in ds.variables.keys():
            return False
    return True

def call(resource=[], variable=None, dimension_map=None, calc=None,
         calc_grouping=None, conform_units_to=None, crs=None, memory_limit=None,  prefix=None,
         regrid_destination=None, regrid_options='bil', level_range=None, cdover='python',
         geom=None, output_format_options=None, search_radius_mult=2.,
         select_nearest=False, select_ugid=None, spatial_wrapping=None,
         t_calendar=None, time_region=None,
         time_range=None, dir_output=None, output_format='nc'):
    '''
    ocgis operation call

    :param resource:
    :param variable: variable in the input file to be picked
    :param dimension_map: dimension map in case of unconventional storage of data
    :param calc: ocgis calc syntax for calculation partion
    :param calc_grouping: time aggregate grouping
    :param conform_units_to:
    :param crs: coordinate reference system
    :param memory_limit: limit the amount of data to be loaded into the memory at once \
        if None (default) free memory is detected by birdhouse
    :param level_range: subset of given levels
    :param prefix: string for the file base name
    :param regrid_destination: file path with netCDF file with grid for output file
    :param geom: name of shapefile stored in birdhouse shape cabinet
    :param output_format_options: output options for netCDF e.g compression level()
    :param regrid_destination: file containing the targed grid (griddes.txt or netCDF file)
    :param regrid_options: methods for regridding:
                          'bil' = Bilinear interpolation
                          'bic' = Bicubic interpolation
                          'dis' = Distance-weighted average remapping
                          'nn' = nearest neighbour
                          'con' = First-order conservative remapping
                          'laf' = largest area fraction reamapping
    :param search_radius_mult: search radius for point geometries. All included gridboxes will be returned
    :param select_nearest: nearest neighbour selection for point geometries
    :param select_ugid: ugid for appropriate polygons
    :param spatial_wrapping: how to handle coordinates in case of subsets, options: None (default), 'wrap', 'unwrap'
    :param time_region: select single month
    :param time_range: sequence of two datetime.datetime objects to mark start and end point
    :param dir_output (default= curdir):
    :param output_format:
    :return: output file path
    '''
    LOGGER.info('Start ocgis module call function')
    from ocgis import OcgOperations, RequestDataset, env, DimensionMap, crs
    from ocgis.util.large_array import compute
    from datetime import datetime as dt
    from datetime import date as dd
    from datetime import time as dt_time
    import uuid

    # prepare the environment
    env.DIR_SHPCABINET = config.shapefiles_path()
    env.OVERWRITE = True
    # env.DIR_OUTPUT = dir_output
    # LOGGER.debug(' **** env.DIR_OUTPUT  = %s ' % env.DIR_OUTPUT)

    if dir_output is None:
        dir_output = abspath(curdir)

    # check time_range format:

    if time_range is not None:
        try:
            LOGGER.debug('time_range type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            LOGGER.debug('time_range= %s , %s ' % (time_range[0], time_range[1]))
            #if type(time_range[0] is 'datetime.date'):
            if (isinstance(time_range[0],dd) and not isinstance(time_range[0],dt)):
                time_range = [dt.combine(time_range[0], dt.min.time()),
                              dt.combine(time_range[1], dt.min.time())]
                # time_range = [dt.combine(time_range[0], dt_time(12,0)),
                #               dt.combine(time_range[1], dt_time(12,0))]
            LOGGER.debug('time_range changed to type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            LOGGER.debug('time_range changed to= %s , %s ' % (time_range[0], time_range[1]))
        except:
            LOGGER.exception('failed to confert data to datetime')

    if spatial_wrapping == 'wrap':
        spatial_reorder = True
    else:
        spatial_reorder = False
    LOGGER.debug('spatial_reorder: %s and spatial_wrapping: %s ' % (spatial_reorder, spatial_wrapping))

    if prefix is None:
        prefix = str(uuid.uuid1())
        env.PREFIX = prefix
    #
    # if output_format_options is False:
    #     output_format_options = None
    # elif output_format_options is True:
    #     output_format_options = {'data_model': 'NETCDF4',  # NETCDF4_CLASSIC
    #                              'variable_kwargs': {'zlib': True, 'complevel': 9}}
    # else:
    if output_format_options is not None:
        LOGGER.info('output_format_options are set to %s ' % (output_format_options))

    if type(resource) != list:
        resource = list([resource])
    # execute ocgis
    LOGGER.info('Execute ocgis module call function')

    # # needed for some AFR-44 data
    # dimension_map = DimensionMap()
    # dimension_map.set_variable('x', 'lon', dimension='rlon')
    # dimension_map.set_variable('y', 'lat', dimension='rlat')
    # dimension_map.set_variable('time', 'time', dimension='time')
    # crs=crs.Spherical()

    try:
        LOGGER.debug('call module curdir = %s ' % abspath(curdir))
        rd = RequestDataset(resource, variable=variable, level_range=level_range,
                            dimension_map=dimension_map, conform_units_to=conform_units_to,
                            time_region=time_region, t_calendar=t_calendar, time_range=time_range) # , crs=crs)

        from ocgis.constants import DimensionMapKey
        rd.dimension_map.set_bounds(DimensionMapKey.TIME, None)

        ops = OcgOperations(dataset=rd,
                            output_format_options=output_format_options,
                            dir_output=dir_output,
                            spatial_wrapping=spatial_wrapping,
                            spatial_reorder=spatial_reorder,
                            # regrid_destination=rd_regrid,
                            # options=options,
                            calc=calc,
                            calc_grouping=calc_grouping,
                            geom=geom,
                            output_format=output_format,
                            prefix=prefix,
                            search_radius_mult=search_radius_mult,
                            select_nearest=select_nearest,
                            select_ugid=select_ugid,
                            add_auxiliary_files=False)
        LOGGER.info('OcgOperations set')
    except:
        LOGGER.exception('failed to setup OcgOperations')
        return None

    try:
        LOGGER.info('ocgis module call as ops.execute()')
        geom_file = ops.execute()
    except:
        LOGGER.exception('failed to execute ocgis operation')
        return None

    #
    # try:
    #     from numpy import sqrt
    #     from flyingpigeon.utils import FreeMemory
    #
    #     if memory_limit is None:
    #         f = FreeMemory()
    #         mem_kb = f.user_free
    #         mem_mb = mem_kb / 1024.
    #         mem_limit = mem_mb / 2.  # set limit to half of the free memory
    #     else:
    #         mem_limit = memory_limit
    #
    #     if mem_limit >= 1024. * 4:
    #         mem_limit = 1024. * 4
    #         # 475.0 MB for openDAP
    #
    #     LOGGER.info('memory_limit = %s Mb' % (mem_limit))
    #
    #     data_kb = ops.get_base_request_size()['total']
    #     data_mb = data_kb / 1024.
    #
    #     # data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
    #     LOGGER.info('data_mb  = %s Mb' % (data_mb))
    #
    #     if data_mb <= mem_limit:  # input is smaler than the half of free memory size
    #         try:
    #             LOGGER.info('ocgis module call as ops.execute()')
    #             geom_file = ops.execute()
    #         except Exception as e:
    #             LOGGER.debug('failed to execute ocgis operation')
    #             raise
    #             return None
    #
    #     else:
    #         ##########################
    #         # calcultion of chunk size
    #         ##########################
    #         try:
    #             size = ops.get_base_request_size()
    #             nb_time_coordinates_rd = size['variables'][variable]['temporal']['shape'][0]
    #             element_in_kb = size['total']/reduce(lambda x, y: x*y, size['variables'][variable]['value']['shape'])
    #             element_in_mb = element_in_kb / 1024.
    #             tile_dim = sqrt(mem_limit/(element_in_mb*nb_time_coordinates_rd))  # maximum chunk size
    #
    #             LOGGER.info('ocgis module call compute with chunks')
    #             if calc is None:
    #                 calc = '%s=%s*1' % (variable, variable)
    #                 LOGGER.info('calc set to = %s ' % calc)
    #             ops = OcgOperations(dataset=rd,
    #                                 output_format_options=output_format_options,
    #                                 dir_output=dir_output,
    #                                 spatial_wrapping=spatial_wrapping,
    #                                 spatial_reorder=spatial_reorder,
    #                                 # regrid_destination=rd_regrid,
    #                                 # options=options,
    #                                 calc=calc,
    #                                 calc_grouping=calc_grouping,
    #                                 geom=geom,
    #                                 output_format=output_format,
    #                                 prefix=prefix,
    #                                 search_radius_mult=search_radius_mult,
    #                                 select_nearest=select_nearest,
    #                                 select_ugid=select_ugid,
    #                                 add_auxiliary_files=False)
    #             geom_file = compute(ops, tile_dimension=int(tile_dim), verbose=True)
    #             print 'ocgis calculated'
    #         except Exception as e:
    #             LOGGER.debug('failed to compute ocgis with chunks')
    #             raise
    #             return None
    #     LOGGER.info('Succeeded with ocgis module call function')
    # except:
    #     LOGGER.exception('failed to compare dataload with free memory, calling as execute instead')

    ############################################
    # remapping according to regrid informations
    ############################################
    if regrid_destination is not None:
        try:
            if (cdover=='system'):
                from os import system
                remap = 'remap%s' % regrid_options
                output = '%s.nc' % uuid.uuid1()
                comcdo = 'cdo -O %s,%s %s %s' % (remap,regrid_destination,geom_file,output)
                system(comcdo)
                # need to substitute by subprocess call
            else:
                from tempfile import mkstemp
                from cdo import Cdo
                cdo = Cdo()
                output = '%s.nc' % uuid.uuid1()
                remap = 'remap%s' % regrid_options
                call = [op for op in dir(cdo) if remap in op]
                cmd = "output = cdo.%s('%s',input='%s', output='%s')" \
                      % (str(call[0]), regrid_destination, geom_file, output)
                exec cmd
        except Exception as e:
            LOGGER.debug('failed to remap')
            raise
            return None
    else:
        output = geom_file

    # try:
    #     from flyingpigeon.utils import unrotate_pole
    #     lat, lon = unrotate_pole(output)
    # except:
    #     LOGGER.exception('failed to unrotate pole')
    return output

def eval_timerange(resource, time_range):
    """
    quality checker if given time_range is covered by timesteps in resource files

    :param resource: input netCDF files
    :param time_range: start and end date of time range [datetime,datetime]

    :returns [datetime,datetime]: time_range
    """
    from flyingpigeon.utils import get_time

    LOGGER.info('time_range: %s' % time_range)

    if type(resource) != str:
        resource.sort()
    time = get_time(resource)
    start = time[0]
    end = time[-1]

    if (time_range[0] > start or time_range[0] < end):
        LOGGER.debug('time range start %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[0] = start
    LOGGER.debug('time_range start changed to first timestep of dataset')
    if (time_range[1] > end or time_range[1] < start):
        LOGGER.debug('time range end %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[1] = end
    LOGGER.debug('time_range end changed to last timestep of dataset')
    if (time_range[0] > time_range[1]):
        time_range = reversed(time_range)
        LOGGER.debug('time range reversed! start was later than end ')
    LOGGER.info('time range start and end set')
    return time_range

# # check memory load
# from os import stat
#   if memory_limit == None:
#     f = FreeMemory()
#     mem_kb = f.user_free
#     mem_mb = mem_kb / 1024.
#     mem_limit = mem_mb / 2. # set limit to half of the free memory
#   else:
#     mem_limit = memory_limit
#
#   if mem_limit >= 1024. * 4:
#     mem_limit = 1024. * 4
#     # 475.0 MB for openDAP
#
#   #if type(resource) == list :
#     #data_kb =  stat(resource[0]).st_size * len(resource)
#   #else:
#     #data_kb =  stat(resource).st_size
#   size = ops.get_base_request_size()['total']
#   data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
#   data_mb = data_kb / 1024.
#
#   if variable == None:
#     variable = rd.variable
#     LOGGER.info('%s as variable dedected' % (variable))
<|endoftext|>


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

import random, os, sys, time, math, pygame
from vector2 import Vector2
from pygame.locals import *

SPACESWARM_VERSION = (0, 2, 1)

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

WINDOWWIDTH, WINDOWHEIGHT = 800, 600
TEXTCOLOR = WHITE
BACKGROUNDCOLOR = BLACK
FPS = 40

random.seed()
pygame.init()
pygame.mouse.set_visible(False) # we blit the mouse instead
screen = pygame.display.set_mode((WINDOWWIDTH, WINDOWHEIGHT))
pygame.display.set_caption('Space Swarm!')

def load_image(name):
    fullname = os.path.join('data', name)
    image = pygame.image.load(fullname)
    if image.get_alpha is None:
        image = image.convert()
    else:
        image = image.convert_alpha()
    return image, image.get_rect()

def load_sound(name):
    class NoneSound:
        def play(self): pass
    if not pygame.mixer: return NoneSound()
    fullname=os.path.join('data', name)
    return pygame.mixer.Sound(fullname)

def terminate():
    pygame.quit()
    sys.exit()

def wait_for_player():
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                terminate()
            if event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    terminate()
                return

def get_n_points_on_circle(center, radius, n=10):
    alpha = math.pi * 2. / n
    points = list(range(n))
    for i in points:
        theta = alpha * i
        point_on_circle = Vector2(math.cos(theta) * radius,
                                  math.sin(theta) * radius)
        points[i] = center + point_on_circle
    return points

def draw_text(text, font, surface, x, y, color=TEXTCOLOR):
    text = font.render(text, 1, color)
    rect = text.get_rect()
    rect.topleft = (x, y)
    surface.blit(text, rect)

class GameObject(pygame.sprite.Sprite):
    def __init__(self, image, rect, destination=None):
        pygame.sprite.Sprite.__init__(self, self.containers)
        self.image = image[0]
        self.rect = rect
        self.destination = destination

    def move(self, time_passed_seconds, speed):
        dv = Vector2(self.destination)
        lv = Vector2(self.rect.x, self.rect.y)
        heading = Vector2.from_points(lv, dv)
        heading.normalize()
        d = heading * time_passed_seconds * speed
        self.rect.move_ip(round(d.x), round(d.y))

class Player(GameObject):
    image = load_image("player.png")

    def __init__(self):
        rect = Player.image[1]
        rect.center = (WINDOWWIDTH / 2, WINDOWHEIGHT / 2)
        GameObject.__init__(self, Player.image, rect)

    def update(self, time_passed):
        pass

class Explosion(GameObject):
    image = load_image("explosion.png")

    def __init__(self, rect):
        GameObject.__init__(self, Explosion.image, rect)
        self._ttl = 5 # number of frames explosion should be visible

    def update(self, time_passed):
        self._ttl -= 1
        if self._ttl == 0: self.kill()

class Alien(GameObject):
    image = load_image("alien.png")
    width, height = image[0].get_size()

    def __init__(self, speed=100, img=None):
        if img is None: img = Alien.image
        GameObject.__init__(self, img,
                            self._random_spawn_rect(),
                            (WINDOWWIDTH/2, WINDOWHEIGHT/2))
        self._speed = speed

    def speed(self):
        """ Gives a slight random variation in speed for every alien """
        return self._speed + random.randint(-5, 5)

    def _random_spawn_rect(self):
        """
        Figures out where the alien should be spawned, which should be somewhere
        along the edges of the screen.
        """
        spawn_loc = random.randint(0, 3)
        x,y = 0,0
        if spawn_loc == 0: # top
            x = random.randint(0, WINDOWWIDTH)
            y = 0
        elif spawn_loc == 1: # right
            x = WINDOWWIDTH
            y = random.randint(0, WINDOWHEIGHT)
        elif spawn_loc == 2: # bottom
            x = random.randint(0, WINDOWWIDTH)
            y = WINDOWHEIGHT
        elif spawn_loc == 3: # left
            x = 0
            y = random.randint(0, WINDOWHEIGHT)

        return pygame.Rect(x, y, Alien.width, Alien.height)

    def update(self, time_passed):
        super(Alien, self).move(time_passed, self.speed())

class SmartAlien(Alien):
    image = load_image("smart_alien.png")
    width, height = image[0].get_size()

    def __init__(self, speed=100):
        Alien.__init__(self, speed, SmartAlien.image)
        self._true_destination = self.destination
        self._new_destination()

    def _new_destination(self):
        lv = Vector2(self.rect.x, self.rect.y)
        candidates = get_n_points_on_circle(lv, 75)
        random.shuffle(candidates)
        for i in candidates:
            # find first point that is closer
            if i.get_distance_to(self._true_destination) < \
                   lv.get_distance_to(self._true_destination):
                self.destination = (i.x, i.y)
                break

    def update(self, time_passed):
        super(Alien, self).update(time_passed)
        lv = Vector2(self.rect.x, self.rect.y)
        dv = Vector2(self.destination)
        if lv.get_distance_to(dv) < 2:
            self._new_destination()

class Bullet(GameObject):
    image = load_image("bullet.png")
    width, height = image[0].get_size()
    speed = 200

    def __init__(self, rect):
        GameObject.__init__(self, Bullet.image,
                            pygame.Rect(WINDOWWIDTH/2, WINDOWHEIGHT/2,
                                        Bullet.width, Bullet.height),
                            self._calculate_destination(rect))

    def _calculate_destination(self, mouse_pos):
        """
        Figure out the destination coords for the bullet, starting from the
        center of the screen through the point the player click, to the edge
        of the screen.
        """
        dx,dy = mouse_pos

        step = Vector2.from_points((WINDOWWIDTH/2,WINDOWHEIGHT/2),
                                   (dx,dy)) * 0.1
        while True:
            if dx > WINDOWWIDTH or dy > WINDOWHEIGHT or dx < 0 or dy < 0:
                break
            dx += step.x
            dy += step.y

        return (dx, dy)

    def update(self, time_passed):
        self.move(time_passed, Bullet.speed)
        if self.rect.top <= 0 or self.rect.bottom >= WINDOWHEIGHT \
               or self.rect.left <= 0 or self.rect.right >= WINDOWWIDTH:
            self.kill()

class Spawner(object):
    def __init__(self, klass, speed, n):
        self.klass = klass
        self.speed = speed
        self.n = n

    def spawn(self):
        self.n -= 1
        return self.klass(self.speed)

# Level definitions
# TODO refactor all the level logic into a class
def instanciate_levels():
    return {
        1: { 'aliens': Spawner(Alien, 50, 4),
         'spawn_rate': 50, 'multiplier': 1 },
        2: { 'aliens': Spawner(Alien, 55, 6),
             'spawn_rate': 45, 'multiplier': 1 },
        3: { 'aliens': Spawner(Alien, 62, 10),
             'spawn_rate': 40, 'multiplier': 1 },
        4: { 'aliens': Spawner(Alien, 67, 14),
             'spawn_rate': 30, 'multiplier': 1 },
        5: { 'aliens': Spawner(Alien, 50, 14),
             'spawn_rate': 25, 'multiplier': 2 },
        6: { 'aliens': Spawner(Alien, 50, 20),
             'spawn_rate': 30, 'multiplier': 2 },
        7: { 'aliens': Spawner(Alien, 30, 26),
             'spawn_rate': 40, 'multiplier': 3 },
        8: { 'aliens': Spawner(SmartAlien, 60, 20),
             'spawn_rate': 30, 'multiplier': 2 }
    }

bg = load_image("bg.jpg")
clock = pygame.time.Clock()

# Fonts
title_font = pygame.font.SysFont(None, 48)
font = pygame.font.SysFont(None, 32)

# Load resources
scope_image = load_image("scope.png")[0]
weapon_sound = load_sound("weapon.wav")
alien_killed_sound = load_sound("alienkilled.wav")
game_over_sound = load_sound("gameover.wav")
levelup_sound = load_sound("levelup.wav")
pygame.mixer.music.load(os.path.join("data", "background.mid"))

# show the "Start" screen
screen.blit(*bg)
draw_text('Space Swarm!', title_font, screen, 20,
         20, RED)
draw_text('To defend Earth, fend off the aliens with your missiles.',
         font, screen, 20, 60)
draw_text('Press any key to start.', font, screen, 20, 90)
draw_text("v"+".".join([str(x) for x in SPACESWARM_VERSION]), font, screen,
          20, WINDOWHEIGHT-40)
pygame.display.update()
wait_for_player()

while True:
    # setup
    game_over, game_finished, muted = False, False, False
    aliens = pygame.sprite.Group()
    bullets = pygame.sprite.Group()
    allsprites = pygame.sprite.RenderUpdates()

    Player.containers = allsprites
    Alien.containers = aliens, allsprites
    Bullet.containers = bullets, allsprites
    Explosion.containers = allsprites

    aliens_killed = 0
    alien_spawn_timer = 0
    score = 0
    level = 1
    levels = instanciate_levels()
    level_dict = levels[level]
    pygame.mixer.music.play(-1, 0.0)

    player = Player()

    while True: # Game loop
        for event in pygame.event.get():
            if event.type is MOUSEBUTTONDOWN: # weapon fired
                if not muted: weapon_sound.play()
                score -= 5 * level # each bullet costs score
                if score < 0: score = 0
                Bullet(pygame.mouse.get_pos())
            elif event.type is KEYDOWN:
                if event.key == K_ESCAPE or event.key == K_q:
                    terminate()
                elif event.key == K_p:
                    wait_for_player()
                elif event.key == K_m:
                    if muted:
                        pygame.mixer.music.unpause()
                        muted = False
                    else:
                        pygame.mixer.music.pause()
                        muted = True
            elif event.type is QUIT:
                terminate()

        alien_spawn_timer += 1
        if alien_spawn_timer == level_dict['spawn_rate'] and \
               level_dict['aliens'].n > 0:
            alien_spawn_timer = 0
            for i in range(level_dict['multiplier']):
                level_dict['aliens'].spawn()
                if level_dict['aliens'].n == 0: break

        # collision detection
        if pygame.sprite.spritecollide(player, aliens, 1):
            game_over = True
            break

        for a in pygame.sprite.groupcollide(aliens, bullets, 1, 1).keys():
            Explosion(a.rect)
            a.kill()
            aliens_killed += 1
            score += 10 * level
            if not muted: alien_killed_sound.play()

        if level_dict['aliens'].n == 0 and len(aliens) == 0:
            if not muted: levelup_sound.play()
            level += 1
            if not level in levels.keys():
                game_finished = True
                break
            else:
                level_dict = levels[level]
                alien_spawn_timer = 0

        # Redraw screen
        screen.blit(*bg)
        draw_text('Level: %s' % (level), font, screen, 0, 0)
        draw_text('Remaining aliens: %s' % (level_dict['aliens'].n),
                 font, screen, 0, 20)
        draw_text('Aliens killed: %s' % (aliens_killed),
                 font, screen, WINDOWWIDTH/2, 0)
        draw_text('Score: %s' % (score), font, screen, WINDOWWIDTH/2, 20)

        allsprites.update(clock.tick(FPS) / 1000.)
        allsprites.draw(screen)
        screen.blit(scope_image, pygame.mouse.get_pos())
        pygame.display.update()

    # broken out of game loop
    pygame.mixer.music.stop()
    if game_over:
        if not muted: game_over_sound.play()
        draw_text('GAME OVER', title_font, screen, (WINDOWWIDTH / 3),
                 (WINDOWHEIGHT / 3), RED)
        draw_text('Press any key to play again, or Esc to quit.', font,
             screen, (WINDOWWIDTH / 3) - 80, (WINDOWHEIGHT / 3) + 50)
    else:
        # TODO game won sound
        draw_text('CONGRATULATIONS!', title_font, screen,
                  (WINDOWWIDTH / 3), (WINDOWHEIGHT / 3), BLUE)
        draw_text('You have saved Earth!', title_font, screen,
                  (WINDOWWIDTH / 3), (WINDOWHEIGHT / 3) + 100, GREEN)
        draw_text('Press any key to play again, or Esc to quit.', font,
             screen, (WINDOWWIDTH / 3) - 80, (WINDOWHEIGHT / 3) + 150)
    pygame.display.update()
    wait_for_player()
<commit_msg>Game now runs even if the pygame mixer isn't initialized<commit_after>#!/usr/bin/env python

import random, os, sys, time, math, pygame
from vector2 import Vector2
from pygame.locals import *

SPACESWARM_VERSION = (0, 2, 1)

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

WINDOWWIDTH, WINDOWHEIGHT = 800, 600
TEXTCOLOR = WHITE
BACKGROUNDCOLOR = BLACK
FPS = 40

random.seed()
pygame.init()
pygame.mouse.set_visible(False) # we blit the mouse instead
screen = pygame.display.set_mode((WINDOWWIDTH, WINDOWHEIGHT))
pygame.display.set_caption('Space Swarm!')

def load_image(name):
    fullname = os.path.join('data', name)
    image = pygame.image.load(fullname)
    if image.get_alpha is None:
        image = image.convert()
    else:
        image = image.convert_alpha()
    return image, image.get_rect()

def load_sound(name):
    class NoneSound:
        def play(self): pass
    if not pygame.mixer.get_init(): return NoneSound()
    fullname=os.path.join('data', name)
    return pygame.mixer.Sound(fullname)

def terminate():
    pygame.quit()
    sys.exit()

def wait_for_player():
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                terminate()
            if event.type == KEYDOWN:
                if event.key == K_ESCAPE:
                    terminate()
                return

def get_n_points_on_circle(center, radius, n=10):
    alpha = math.pi * 2. / n
    points = list(range(n))
    for i in points:
        theta = alpha * i
        point_on_circle = Vector2(math.cos(theta) * radius,
                                  math.sin(theta) * radius)
        points[i] = center + point_on_circle
    return points

def draw_text(text, font, surface, x, y, color=TEXTCOLOR):
    text = font.render(text, 1, color)
    rect = text.get_rect()
    rect.topleft = (x, y)
    surface.blit(text, rect)

class GameObject(pygame.sprite.Sprite):
    def __init__(self, image, rect, destination=None):
        pygame.sprite.Sprite.__init__(self, self.containers)
        self.image = image[0]
        self.rect = rect
        self.destination = destination

    def move(self, time_passed_seconds, speed):
        dv = Vector2(self.destination)
        lv = Vector2(self.rect.x, self.rect.y)
        heading = Vector2.from_points(lv, dv)
        heading.normalize()
        d = heading * time_passed_seconds * speed
        self.rect.move_ip(round(d.x), round(d.y))

class Player(GameObject):
    image = load_image("player.png")

    def __init__(self):
        rect = Player.image[1]
        rect.center = (WINDOWWIDTH / 2, WINDOWHEIGHT / 2)
        GameObject.__init__(self, Player.image, rect)

    def update(self, time_passed):
        pass

class Explosion(GameObject):
    image = load_image("explosion.png")

    def __init__(self, rect):
        GameObject.__init__(self, Explosion.image, rect)
        self._ttl = 5 # number of frames explosion should be visible

    def update(self, time_passed):
        self._ttl -= 1
        if self._ttl == 0: self.kill()

class Alien(GameObject):
    image = load_image("alien.png")
    width, height = image[0].get_size()

    def __init__(self, speed=100, img=None):
        if img is None: img = Alien.image
        GameObject.__init__(self, img,
                            self._random_spawn_rect(),
                            (WINDOWWIDTH/2, WINDOWHEIGHT/2))
        self._speed = speed

    def speed(self):
        """ Gives a slight random variation in speed for every alien """
        return self._speed + random.randint(-5, 5)

    def _random_spawn_rect(self):
        """
        Figures out where the alien should be spawned, which should be somewhere
        along the edges of the screen.
        """
        spawn_loc = random.randint(0, 3)
        x,y = 0,0
        if spawn_loc == 0: # top
            x = random.randint(0, WINDOWWIDTH)
            y = 0
        elif spawn_loc == 1: # right
            x = WINDOWWIDTH
            y = random.randint(0, WINDOWHEIGHT)
        elif spawn_loc == 2: # bottom
            x = random.randint(0, WINDOWWIDTH)
            y = WINDOWHEIGHT
        elif spawn_loc == 3: # left
            x = 0
            y = random.randint(0, WINDOWHEIGHT)

        return pygame.Rect(x, y, Alien.width, Alien.height)

    def update(self, time_passed):
        super(Alien, self).move(time_passed, self.speed())

class SmartAlien(Alien):
    image = load_image("smart_alien.png")
    width, height = image[0].get_size()

    def __init__(self, speed=100):
        Alien.__init__(self, speed, SmartAlien.image)
        self._true_destination = self.destination
        self._new_destination()

    def _new_destination(self):
        lv = Vector2(self.rect.x, self.rect.y)
        candidates = get_n_points_on_circle(lv, 75)
        random.shuffle(candidates)
        for i in candidates:
            # find first point that is closer
            if i.get_distance_to(self._true_destination) < \
                   lv.get_distance_to(self._true_destination):
                self.destination = (i.x, i.y)
                break

    def update(self, time_passed):
        super(Alien, self).update(time_passed)
        lv = Vector2(self.rect.x, self.rect.y)
        dv = Vector2(self.destination)
        if lv.get_distance_to(dv) < 2:
            self._new_destination()

class Bullet(GameObject):
    image = load_image("bullet.png")
    width, height = image[0].get_size()
    speed = 200

    def __init__(self, rect):
        GameObject.__init__(self, Bullet.image,
                            pygame.Rect(WINDOWWIDTH/2, WINDOWHEIGHT/2,
                                        Bullet.width, Bullet.height),
                            self._calculate_destination(rect))

    def _calculate_destination(self, mouse_pos):
        """
        Figure out the destination coords for the bullet, starting from the
        center of the screen through the point the player click, to the edge
        of the screen.
        """
        dx,dy = mouse_pos

        step = Vector2.from_points((WINDOWWIDTH/2,WINDOWHEIGHT/2),
                                   (dx,dy)) * 0.1
        while True:
            if dx > WINDOWWIDTH or dy > WINDOWHEIGHT or dx < 0 or dy < 0:
                break
            dx += step.x
            dy += step.y

        return (dx, dy)

    def update(self, time_passed):
        self.move(time_passed, Bullet.speed)
        if self.rect.top <= 0 or self.rect.bottom >= WINDOWHEIGHT \
               or self.rect.left <= 0 or self.rect.right >= WINDOWWIDTH:
            self.kill()

class Spawner(object):
    def __init__(self, klass, speed, n):
        self.klass = klass
        self.speed = speed
        self.n = n

    def spawn(self):
        self.n -= 1
        return self.klass(self.speed)

# Level definitions
# TODO refactor all the level logic into a class
def instanciate_levels():
    return {
        1: { 'aliens': Spawner(Alien, 50, 4),
         'spawn_rate': 50, 'multiplier': 1 },
        2: { 'aliens': Spawner(Alien, 55, 6),
             'spawn_rate': 45, 'multiplier': 1 },
        3: { 'aliens': Spawner(Alien, 62, 10),
             'spawn_rate': 40, 'multiplier': 1 },
        4: { 'aliens': Spawner(Alien, 67, 14),
             'spawn_rate': 30, 'multiplier': 1 },
        5: { 'aliens': Spawner(Alien, 50, 14),
             'spawn_rate': 25, 'multiplier': 2 },
        6: { 'aliens': Spawner(Alien, 50, 20),
             'spawn_rate': 30, 'multiplier': 2 },
        7: { 'aliens': Spawner(Alien, 30, 26),
             'spawn_rate': 40, 'multiplier': 3 },
        8: { 'aliens': Spawner(SmartAlien, 60, 20),
             'spawn_rate': 30, 'multiplier': 2 }
    }

bg = load_image("bg.jpg")
clock = pygame.time.Clock()

# Fonts
title_font = pygame.font.SysFont(None, 48)
font = pygame.font.SysFont(None, 32)

# Load resources
scope_image = load_image("scope.png")[0]
weapon_sound = load_sound("weapon.wav")
alien_killed_sound = load_sound("alienkilled.wav")
game_over_sound = load_sound("gameover.wav")
levelup_sound = load_sound("levelup.wav")
if pygame.mixer.get_init():
    pygame.mixer.music.load(os.path.join("data", "background.mid"))

# show the "Start" screen
screen.blit(*bg)
draw_text('Space Swarm!', title_font, screen, 20,
         20, RED)
draw_text('To defend Earth, fend off the aliens with your missiles.',
         font, screen, 20, 60)
draw_text('Press any key to start.', font, screen, 20, 90)
draw_text("v"+".".join([str(x) for x in SPACESWARM_VERSION]), font, screen,
          20, WINDOWHEIGHT-40)
pygame.display.update()
wait_for_player()

while True:
    # setup
    game_over, game_finished, muted = False, False, False
    aliens = pygame.sprite.Group()
    bullets = pygame.sprite.Group()
    allsprites = pygame.sprite.RenderUpdates()

    Player.containers = allsprites
    Alien.containers = aliens, allsprites
    Bullet.containers = bullets, allsprites
    Explosion.containers = allsprites

    aliens_killed = 0
    alien_spawn_timer = 0
    score = 0
    level = 1
    levels = instanciate_levels()
    level_dict = levels[level]
    if pygame.mixer.get_init(): pygame.mixer.music.play(-1, 0.0)

    player = Player()

    while True: # Game loop
        for event in pygame.event.get():
            if event.type is MOUSEBUTTONDOWN: # weapon fired
                if not muted: weapon_sound.play()
                score -= 5 * level # each bullet costs score
                if score < 0: score = 0
                Bullet(pygame.mouse.get_pos())
            elif event.type is KEYDOWN:
                if event.key == K_ESCAPE or event.key == K_q:
                    terminate()
                elif event.key == K_p:
                    wait_for_player()
                elif event.key == K_m:
                    if muted:
                        pygame.mixer.music.unpause()
                        muted = False
                    else:
                        pygame.mixer.music.pause()
                        muted = True
            elif event.type is QUIT:
                terminate()

        alien_spawn_timer += 1
        if alien_spawn_timer == level_dict['spawn_rate'] and \
               level_dict['aliens'].n > 0:
            alien_spawn_timer = 0
            for i in range(level_dict['multiplier']):
                level_dict['aliens'].spawn()
                if level_dict['aliens'].n == 0: break

        # collision detection
        if pygame.sprite.spritecollide(player, aliens, 1):
            game_over = True
            break

        for a in pygame.sprite.groupcollide(aliens, bullets, 1, 1).keys():
            Explosion(a.rect)
            a.kill()
            aliens_killed += 1
            score += 10 * level
            if not muted: alien_killed_sound.play()

        if level_dict['aliens'].n == 0 and len(aliens) == 0:
            if not muted: levelup_sound.play()
            level += 1
            if not level in levels.keys():
                game_finished = True
                break
            else:
                level_dict = levels[level]
                alien_spawn_timer = 0

        # Redraw screen
        screen.blit(*bg)
        draw_text('Level: %s' % (level), font, screen, 0, 0)
        draw_text('Remaining aliens: %s' % (level_dict['aliens'].n),
                 font, screen, 0, 20)
        draw_text('Aliens killed: %s' % (aliens_killed),
                 font, screen, WINDOWWIDTH/2, 0)
        draw_text('Score: %s' % (score), font, screen, WINDOWWIDTH/2, 20)

        allsprites.update(clock.tick(FPS) / 1000.)
        allsprites.draw(screen)
        screen.blit(scope_image, pygame.mouse.get_pos())
        pygame.display.update()

    # broken out of game loop
    pygame.mixer.music.stop()
    if game_over:
        if not muted: game_over_sound.play()
        draw_text('GAME OVER', title_font, screen, (WINDOWWIDTH / 3),
                 (WINDOWHEIGHT / 3), RED)
        draw_text('Press any key to play again, or Esc to quit.', font,
             screen, (WINDOWWIDTH / 3) - 80, (WINDOWHEIGHT / 3) + 50)
    else:
        # TODO game won sound
        draw_text('CONGRATULATIONS!', title_font, screen,
                  (WINDOWWIDTH / 3), (WINDOWHEIGHT / 3), BLUE)
        draw_text('You have saved Earth!', title_font, screen,
                  (WINDOWWIDTH / 3), (WINDOWHEIGHT / 3) + 100, GREEN)
        draw_text('Press any key to play again, or Esc to quit.', font,
             screen, (WINDOWWIDTH / 3) - 80, (WINDOWHEIGHT / 3) + 150)
    pygame.display.update()
    wait_for_player()
<|endoftext|>


<commit_before>"""Python script to do FIO test for GCSFuse Vs GCSFuse/Goofys.

To run the script:
>> python3 compare_fuse_types_using_fio.py -- fuse_type_1 fuse_type_1_version fuser_type_2 fuse_type_2_version jobfile_path

-> Supported fuse_types are gcsfuse and goofys.
-> Incase of gcsfuse you can specify version as any of its released version or master to run from the source.
-> Incase of goofys specify the version as latest, as we will be comparing gcsfuse with the latest version of goofys.

The code takes input the fuse type and its version and, jobfile_path as input for fio test.
Then it will perform the fio test for each of the given fuse type at a specific version
and extract fio_metrics to output.txt file.
"""
import argparse
import os
from fio import fio_metrics

from absl import app

GCS_BUCKET = 'gke_load_test'
GOOFYS_REPO = 'https://github.com/kahing/goofys'

def _install_gcsfuse(version) -> None:
  """Install gcsfuse with Specific version.
  Args:
    version(str): gcsfuse version to be installed.
  """
  os.system(f'''curl -L -O https://github.com/GoogleCloudPlatform/gcsfuse/releases/download/v{version}/gcsfuse_{version}_amd64.deb
            sudo dpkg --install gcsfuse_{version}_amd64.deb
            mkdir gcs
            gcsfuse --implicit-dirs --max-conns-per-host 100 --disable-http2 {GCS_BUCKET} gcs
            ''')

def _remove_gcsfuse(version) -> None:
  """Remove gcsfuse with specific version.
  Args:
    version(str): gcsfuse version to be removed.
  """
  os.system(f'''fusermount -u gcs
            rm -rf gcs
            rm -rf gcsfuse_{version}_amd64.deb
            sudo apt-get remove gcsfuse -y
            ''')

def _install_goofys() -> None:
  """Install latest version of goofys.
  """
  os.system(f'''git clone {GOOFYS_REPO}
            export GOPATH=$HOME/work
            mkdir gcs
            cd goofys
            go run . gs://{GCS_BUCKET} ../gcs
            cd ..
            ''')

def _remove_goofys() -> None:
  """Remove goofys.
  """
  os.system('''fusermount -u gcs
            rm -rf gcs
            rm -rf goofys
            ''')

def _run_fio_test(jobfile_path, fio_metrics_obj) -> None:
  """Run fio test and extract metrics to output.txt file.
  Args:
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  os.system(f'''fio {jobfile_path} --lat_percentiles 1 --output-format=json --output='output.json'
            ''')
  fio_metricss = fio_metrics_obj.get_metrics('output.json', False)
  os.system(f'''echo {fio_metricss} >> out/output.txt
            rm output.json
            ''')

def _gcsfuse_test(version, jobfile_path, fio_metrics_obj) -> None:
  """FIO test for gcsfuse of given version.
  Args:
    version(str): gcsfuse version to perform fio test.
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  _install_gcsfuse(version)
  _run_fio_test(jobfile_path, fio_metrics_obj)
  _remove_gcsfuse(version)

def _goofys_test(jobfile_path, fio_metrics_obj) -> None:
  """FIO test for latest version of goofys.

  Args:
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  _install_goofys()
  _run_fio_test(jobfile_path, fio_metrics_obj)
  _remove_goofys()

def main(argv) -> None:

  parser = argparse.ArgumentParser()
  parser.add_argument(
      'fuse_type_1',
      help='Provide reading fuse_type_1 for fio test')
  parser.add_argument(
      'fuse_type_1_version',
      help='Provid the Specific version of fuse_type_1 in approprite formate')
  parser.add_argument(
      'fuse_type_2',
      help='Provide reading fuse_type_2 for fio test')
  parser.add_argument(
      'fuse_type_2_version',
      help='Provid the Specific version of fuse_type_2 in approprite formate')
  parser.add_argument(
      'jobfile_path',
      help='Provid path of the jobfile')
  args = parser.parse_args(argv[1:])

  fio_metrics_obj = fio_metrics.FioMetrics()
  os.system('mkdir out')

  os.system(f'echo Fuse Type 1: {args.fuse_type_1} {args.fuse_type_1_version} >> out/output.txt')
  if args.method1 == 'gcsfuse':
    _gcsfuse_test(args.fuse_type_1_version, args.jobfile_path, fio_metrics_obj)
  else:
    _goofys_test(args.jobfile_path, fio_metrics_obj)

  os.system(f'echo Fuse Type 2: {args.fuse_type_2} {args.fuse_type_2_version} >> out/output.txt')
  if args.method2 == 'gcsfuse':
    _gcsfuse_test(args.fuse_type_2_version, args.jobfile_path, fio_metrics_obj)
  else:
    _goofys_test(args.jobfile_path, fio_metrics_obj)

if __name__ == '__main__':
  app.run(main)

<commit_msg>Extract fuse test logic into separate function<commit_after>"""Python script to do FIO test for GCSFuse Vs GCSFuse/Goofys.

To run the script:
>> python3 compare_fuse_types_using_fio.py -- fuse_type_1 fuse_type_1_version fuser_type_2 fuse_type_2_version jobfile_path

-> Supported fuse_types are gcsfuse and goofys.
-> Incase of gcsfuse you can specify version as any of its released version or master to run from the source.
-> Incase of goofys specify the version as latest, as we will be comparing gcsfuse with the latest version of goofys.

The code takes input the fuse type and its version and, jobfile_path as input for fio test.
Then it will perform the fio test for each of the given fuse type at a specific version
and extract fio_metrics to output.txt file.
"""
import argparse
import os
from fio import fio_metrics

from absl import app

GCS_BUCKET = 'gke_load_test'
GOOFYS_REPO = 'https://github.com/kahing/goofys'

def _install_gcsfuse(version) -> None:
  """Install gcsfuse with Specific version.
  Args:
    version(str): gcsfuse version to be installed.
  """
  os.system(f'''curl -L -O https://github.com/GoogleCloudPlatform/gcsfuse/releases/download/v{version}/gcsfuse_{version}_amd64.deb
            sudo dpkg --install gcsfuse_{version}_amd64.deb
            mkdir gcs
            gcsfuse --implicit-dirs --max-conns-per-host 100 --disable-http2 {GCS_BUCKET} gcs
            ''')

def _remove_gcsfuse(version) -> None:
  """Remove gcsfuse with specific version.
  Args:
    version(str): gcsfuse version to be removed.
  """
  os.system(f'''fusermount -u gcs
            rm -rf gcs
            rm -rf gcsfuse_{version}_amd64.deb
            sudo apt-get remove gcsfuse -y
            ''')

def _install_goofys() -> None:
  """Install latest version of goofys.
  """
  os.system(f'''git clone {GOOFYS_REPO}
            export GOPATH=$HOME/work
            mkdir gcs
            cd goofys
            go run . gs://{GCS_BUCKET} ../gcs
            cd ..
            ''')

def _remove_goofys() -> None:
  """Remove goofys.
  """
  os.system('''fusermount -u gcs
            rm -rf gcs
            rm -rf goofys
            ''')

def _run_fio_test(jobfile_path, fio_metrics_obj) -> None:
  """Run fio test and extract metrics to output.txt file.
  Args:
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  os.system(f'''fio {jobfile_path} --lat_percentiles 1 --output-format=json --output='output.json'
            ''')
  fio_metricss = fio_metrics_obj.get_metrics('output.json', False)
  os.system(f'''echo {fio_metricss} >> out/output.txt
            rm output.json
            ''')

def _gcsfuse_test(version, jobfile_path, fio_metrics_obj) -> None:
  """FIO test for gcsfuse of given version.
  Args:
    version(str): gcsfuse version to perform fio test.
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  _install_gcsfuse(version)
  _run_fio_test(jobfile_path, fio_metrics_obj)
  _remove_gcsfuse(version)

def _goofys_test(jobfile_path, fio_metrics_obj) -> None:
  """FIO test for latest version of goofys.

  Args:
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  _install_goofys()
  _run_fio_test(jobfile_path, fio_metrics_obj)
  _remove_goofys()

 def _fuse_test(fuse_type, fuse_type_version, jobfile_path, fio_metric_obj) -> None:
  """FIO test for specific version of given fuse type.

  Args:
    fuse_type(str): fuse type for fio test.
    fuse_type_version(str): fuse type version for fio test.
    jobfile(str): path of the job file.
    fio_metrics_obj(str): object for extracting fio metrics.
  """
  switch(fuse_type):
    case 'gcsfuse':
      _gcsfuse_test(fuse_type_version, jobfile_path, fio_metric_obj)
    case 'goofys':
      _goofys_test(jobfile_path, fio_metric_obj)
    default:
      app.UsageError('Unsupported fuse type!')

def main(argv) -> None:

  parser = argparse.ArgumentParser()
  parser.add_argument(
      'fuse_type_1',
      help='Provide reading fuse_type_1 for fio test')
  parser.add_argument(
      'fuse_type_1_version',
      help='Provid the Specific version of fuse_type_1 in approprite formate')
  parser.add_argument(
      'fuse_type_2',
      help='Provide reading fuse_type_2 for fio test')
  parser.add_argument(
      'fuse_type_2_version',
      help='Provid the Specific version of fuse_type_2 in approprite formate')
  parser.add_argument(
      'jobfile_path',
      help='Provid path of the jobfile')
  args = parser.parse_args(argv[1:])

  fio_metrics_obj = fio_metrics.FioMetrics()
  os.system('mkdir out')

  os.system(f'echo Fuse Type 1: {args.fuse_type_1} {args.fuse_type_1_version} >> out/output.txt')
  _fuse_test(args.fuse_type_1, args.fuse_type_1_version, args.jobfile_path, fio_metrics_obj)

  os.system(f'echo Fuse Type 2: {args.fuse_type_2} {args.fuse_type_2_version} >> out/output.txt')
  _fuse_test(args.fuse_type_2, args.fuse_type_2_version, args.jobfile_path, fio_metrics_obj)

if __name__ == '__main__':
  app.run(main)

<|endoftext|>


<commit_before>from rivr.http import ResponseRedirect
from rivr_jinja import JinjaView, JinjaResponse

from sotu.github import *
from sotu.models import Entrant

class IndexView(JinjaView):
    template_name = 'index.html'

    def get_context_data(self, **kwargs):
        parameters = {
            'client_id': GITHUB_CLIENT_ID,
            'redirect_uri': GITHUB_CALLBACK_URI,
            'scope': 'user:email',
            'state': 'unused',
        }

        return {
            'github_authorize_uri': 'https://github.com/login/oauth/authorize?' + urllib.urlencode(parameters)
        }

class EntrantView(JinjaView):
    template_name = 'entrant.html'

    def get_context_data(self, **kwargs):
        return {
            'entrant': self.entrant,
            'avatar': self.avatar,
        }

def callback(request):
    code = request.GET.get('code')
    if not code:
        error = request.GET['error']
        error_description = request.GET['error_description']
        return JinjaResponse(request, template_names=['error.html'],
                context={'reason': error_description})

    access_token = retrieve_access_token(code)
    if not access_token:
        return ResponseRedirect('https://sotu.cocoapods.org/')
    user = retrieve_account(access_token)
    email = retrieve_email(access_token)
    username = user['login']

    name = user.get('name', username)
    if len(name) == 0:
        name = username

    avatar = user.get('avatar_url', None)

    try:
        entrant = Entrant.select().where(Entrant.github_username == username).get()
    except Entrant.DoesNotExist:
        entrant = Entrant.create(github_username=username, name=name, email=email)

    return EntrantView.as_view(entrant=entrant, avatar=avatar)(request)

<commit_msg>Handle the lack of error or success from github redirect URI<commit_after>from rivr.http import ResponseRedirect
from rivr_jinja import JinjaView, JinjaResponse

from sotu.github import *
from sotu.models import Entrant

class IndexView(JinjaView):
    template_name = 'index.html'

    def get_context_data(self, **kwargs):
        parameters = {
            'client_id': GITHUB_CLIENT_ID,
            'redirect_uri': GITHUB_CALLBACK_URI,
            'scope': 'user:email',
            'state': 'unused',
        }

        return {
            'github_authorize_uri': 'https://github.com/login/oauth/authorize?' + urllib.urlencode(parameters)
        }

class EntrantView(JinjaView):
    template_name = 'entrant.html'

    def get_context_data(self, **kwargs):
        return {
            'entrant': self.entrant,
            'avatar': self.avatar,
        }

def callback(request):
    code = request.GET.get('code')
    if not code:
        error = request.GET.get('error')
        error_description = request.GET.get('error_description', 'Unknown Error, please try again')
        return JinjaResponse(request, template_names=['error.html'],
                context={'reason': '{} ({})'.format(error_description, error)})

    access_token = retrieve_access_token(code)
    if not access_token:
        return ResponseRedirect('https://sotu.cocoapods.org/')
    user = retrieve_account(access_token)
    email = retrieve_email(access_token)
    username = user['login']

    name = user.get('name', username)
    if len(name) == 0:
        name = username

    avatar = user.get('avatar_url', None)

    try:
        entrant = Entrant.select().where(Entrant.github_username == username).get()
    except Entrant.DoesNotExist:
        entrant = Entrant.create(github_username=username, name=name, email=email)

    return EntrantView.as_view(entrant=entrant, avatar=avatar)(request)

<|endoftext|>


<commit_before>"""
The main entry point for pakit.

Acts as an intermediary between program arguments and pakit Tasks.
"""
from __future__ import absolute_import, print_function

import argparse
import logging
import logging.handlers
import os
import sys
import traceback

from pakit import __version__
from pakit.conf import Config, InstallDB
from pakit.exc import PakitError, PakitDBError
from pakit.recipe import RecipeDB
from pakit.task import (
    InstallTask, RemoveTask, UpdateTask, ListInstalled, ListAvailable,
    DisplayTask, SearchTask
)
import pakit.conf
import pakit.shell

PLOG = logging.getLogger('pakit')

def args_parser():
    """
    Create the program argument parser.

    Returns:
        An argparse parser object.
    """
    prog_name = os.path.basename(os.path.dirname(sys.argv[0]))
    mesg = """
    {0} is a build tool providing a package manager like interface
    to build & install recipes into local paths.

    First Time Run:
        Run in shell: `{0} --create-conf`
            This writes the default config to $HOME/.pakit.yml

        Run in shell: `export PATH=/tmp/pakit/links/bin:$PATH`
            This is where all programs eventually linked to by default.
            See man page for more configuration information.

    For project development or to report bugs:
        https://github.com/starcraftman/pakit

    Additional Information:
        - Man page, provided inside `pakit/extra/pakit.1`. It is automatically
        available via `man pakit` if `paths.link` is on your shell $PATH.
        - `pydoc pakit`
        - See DEMO.md inside pakit site package.

    Recipes:
        pakit_recipes/ag.py is an example recipe
        User made recipes can be put in $HOME/.pakit/recipes by default.
        If two recipes have same name, last one in will be executed.
    """.format(prog_name)
    parser = argparse.ArgumentParser(prog=prog_name, description=mesg,
                                     formatter_class=argparse.
                                     RawDescriptionHelpFormatter)
    mut1 = parser.add_mutually_exclusive_group()
    parser.add_argument('-v', '--version', action='version',
                        version='pakit {0}\nALPHA!'.format(__version__))
    parser.add_argument('-a', '--available', default=False,
                        action='store_true', help='list available recipes')
    parser.add_argument('--available-short', default=False,
                        action='store_true',
                        help='list available recipes, terse output')
    parser.add_argument('-c', '--conf',
                        default=os.path.expanduser('~/.pakit.yml'),
                        help='yaml config file')
    parser.add_argument('--create-conf', default=False, action='store_true',
                        help='write the default config to CONF')
    parser.add_argument('-d', '--display', nargs='+', metavar='PROG',
                        help='show detailed information on recipe')
    mut1.add_argument('-i', '--install', nargs='+',
                      metavar='PROG', help='install specified program(s)')
    mut1.add_argument('-k', '--search', nargs='+', metavar='WORD',
                      help='search names & descriptions for WORD')
    parser.add_argument('-l', '--list', default=False,
                        action='store_true', help='list installed programs')
    parser.add_argument('--list-short', default=False, action='store_true',
                        help='list installed recipes, terse output')
    mut1.add_argument('-r', '--remove', nargs='+',
                      metavar='PROG', help='remove specified program(s)')
    mut1.add_argument('-u', '--update', default=False, action='store_true',
                      help='update installed programs')

    return parser

def global_init(config_file):
    """
    Performs global configuration of pakit.

    Must be called before using pakit. Will ...
        - Read user configuration.
        - Initialize the logging system.
        - Populate the recipe database.
        - Create configured folders.
        - Setup pakit man page.

    Args:
        config_file: The YAML configuration filename.

    Returns:
        The loaded config object.
    """
    config = Config(search_for_config(config_file))
    pakit.conf.CONFIG = config
    log_init(config)
    logging.debug('Global Config: %s', config)

    for path in config.get('pakit.paths').values():
        try:
            os.makedirs(path)
        except OSError:
            pass

    prefix = config.get('pakit.paths.prefix')
    pakit.conf.IDB = InstallDB(os.path.join(prefix, 'installed.yaml'))
    logging.debug('InstallDB: %s', pakit.conf.IDB)
    pakit.shell.TMP_DIR = os.path.dirname(prefix)

    recipes = RecipeDB(config)
    default = os.path.join(os.path.dirname(os.path.dirname(__file__)),
                           'pakit_recipes')
    recipes.index(default)
    try:
        recipes.index(config.get('pakit.paths.recipes'))
    except KeyError:
        pass

    link_man_page(config.get('pakit.paths.link'))

    return config

def link_man_page(link_dir):
    """
    Silently links man page into link dir.
    """
    src = os.path.join(os.path.dirname(__file__), 'extra', 'pakit.1')
    dst = os.path.join(link_dir, 'share', 'man', 'man1', 'pakit.1')
    try:
        os.makedirs(os.path.dirname(dst))
    except OSError:
        pass
    try:
        os.symlink(src, dst)
    except OSError:
        pass

def log_init(config):
    """
    Setup project wide logging.

    Specifically this will setup both file and console logging.
    """
    log_file = config.get('pakit.log.file')
    try:
        os.makedirs(os.path.dirname(log_file))
    except OSError:
        pass

    root = logging.getLogger()
    root.setLevel(logging.DEBUG)
    log_fmt = '%(levelname)s %(asctime)s %(threadName)s ' \
              '%(filename)s %(message)s'
    my_fmt = logging.Formatter(fmt=log_fmt, datefmt='[%d/%m %H%M.%S]')

    while len(root.handlers):
        root.removeHandler(root.handlers[0])

    rot = logging.handlers.RotatingFileHandler(log_file,
                                               mode='a',
                                               maxBytes=2 ** 22,
                                               backupCount=4)
    rot.setLevel(logging.DEBUG)
    rot.setFormatter(my_fmt)
    root.addHandler(rot)

    stream = logging.StreamHandler()
    stream.setLevel(logging.ERROR)
    stream.setFormatter(my_fmt)
    root.addHandler(stream)

    # Logger for informing user
    pak = logging.getLogger('pakit')
    while len(pak.handlers):
        pak.removeHandler(pak.handlers[0])

    pak.setLevel(logging.INFO)
    pak_fmt = 'pakit: %(asctime)s %(message)s'
    pak_info = logging.Formatter(fmt=pak_fmt, datefmt='[%H:%M:%S]')
    pak_stream = logging.StreamHandler()
    pak_stream.setFormatter(pak_info)
    pak.addHandler(pak_stream)

def parse_tasks(args):
    """
    Parse the program arguments into a list of Tasks to execute

    Args:
        args: An argparse object to map onto tasks.

    Returns:
        A list of Tasks.
    """
    tasks = []

    if args.install:
        tasks.extend([InstallTask(prog) for prog in args.install])
    if args.remove:
        tasks.extend([RemoveTask(prog) for prog in args.remove])
    if args.update:
        tasks.extend([UpdateTask(prog) for prog, _ in pakit.conf.IDB])
    if args.available:
        tasks.append(ListAvailable(False))
    if args.available_short:
        tasks.append(ListAvailable(True))
    if args.display:
        tasks.extend([DisplayTask(prog) for prog in args.display])
    if args.search:
        tasks.append(SearchTask(RecipeDB().names(desc=True), args.search))
    if args.list:
        tasks.append(ListInstalled(False))
    if args.list_short:
        tasks.append(ListInstalled(True))

    return tasks

def search_for_config(default_config=None):
    """
    Search for the most relevant configuration file.

    Will search from current dir upwards until root for config files matching:
        - .pakit.yml
        - .pakit.yaml
        - pakit.yml
        - pakit.yaml

    If nothing found, search in $HOME and then $HOME/.pakit for same files.

    If still nothing found, return default_config.

    Args:
        default_config: The default config if nothing else present.

    Returns:
        If a path existed, return found filename.
        If no paths matched an existing file return the default_config.
    """
    folders = [os.getcwd()]
    cur_dir = os.path.dirname(os.getcwd())
    while cur_dir != '/':
        folders.append(cur_dir)
        cur_dir = os.path.dirname(cur_dir)
    folders.append(os.path.expanduser('~'))
    folders.append(os.path.expanduser('~/.pakit'))

    for folder in folders:
        for conf in ['.pakit.yml', '.pakit.yaml', 'pakit.yml', 'pakit.yaml']:
            config_file = os.path.join(folder, conf)
            if os.path.exists(config_file):
                return config_file

    return default_config

def write_config(config_file):
    """
    Writes the DEFAULT config to the config file.
    Overwrites the file if present.

    Raises:
        PakitError: File exists and is a directory.
        PakitError: File could not be written to.
    """
    config = Config(config_file)
    try:
        os.remove(config.filename)
    except OSError:
        if os.path.isdir(config.filename):
            raise PakitError('Config path is a directory.')
    try:
        config.reset()
        config.write()
        print('pakit: Wrote default config to: ' + config.filename)
    except (IOError, OSError):
        raise PakitError('Failed to write to ' + config.filename)
    finally:
        sys.exit(0)

# TODO: Path modification during operation by os.environ
def main(argv=None):
    """
    The main entry point for this program.

    Args:
        argv: A list of program options, if None use sys.argv.
    """
    if argv is None:
        argv = sys.argv

    parser = args_parser()
    if len(argv) == 1:
        logging.error('No arguments. What should I do?')
        parser.print_usage()
        sys.exit(1)

    args = parser.parse_args(argv[1:])
    if args.create_conf:
        write_config(args.conf)

    config = global_init(args.conf)
    logging.debug('CLI: %s', args)
    PLOG.info('Loading config from: ' + config.filename)

    try:
        tasks = parse_tasks(args)
        for task in tasks:
            PLOG.info('Running: %s', str(task))
            task.run()

        if len(tasks) == 0 and args.update:
            PLOG.info('Nothing to update.')

    except PakitDBError as exc:
        PLOG.info(str(exc))
    except PakitError as exc:
        logging.error(exc)
        logging.error(traceback.format_exc())
        raise

if __name__ == '__main__':
    main()  # pragma: no cover
<commit_msg>Logically, if specified --conf should always be used.<commit_after>"""
The main entry point for pakit.

Acts as an intermediary between program arguments and pakit Tasks.
"""
from __future__ import absolute_import, print_function

import argparse
import logging
import logging.handlers
import os
import sys
import traceback

from pakit import __version__
from pakit.conf import Config, InstallDB
from pakit.exc import PakitError, PakitDBError
from pakit.recipe import RecipeDB
from pakit.task import (
    InstallTask, RemoveTask, UpdateTask, ListInstalled, ListAvailable,
    DisplayTask, SearchTask
)
import pakit.conf
import pakit.shell

PLOG = logging.getLogger('pakit')

def args_parser():
    """
    Create the program argument parser.

    Returns:
        An argparse parser object.
    """
    prog_name = os.path.basename(os.path.dirname(sys.argv[0]))
    mesg = """
    {0} is a build tool providing a package manager like interface
    to build & install recipes into local paths.

    First Time Run:
        Run in shell: `{0} --create-conf`
            This writes the default config to $HOME/.pakit.yml

        Run in shell: `export PATH=/tmp/pakit/links/bin:$PATH`
            This is where all programs eventually linked to by default.
            See man page for more configuration information.

    For project development or to report bugs:
        https://github.com/starcraftman/pakit

    Additional Information:
        - Man page, provided inside `pakit/extra/pakit.1`. It is automatically
        available via `man pakit` if `paths.link` is on your shell $PATH.
        - `pydoc pakit`
        - See DEMO.md inside pakit site package.

    Recipes:
        pakit_recipes/ag.py is an example recipe
        User made recipes can be put in $HOME/.pakit/recipes by default.
        If two recipes have same name, last one in will be executed.
    """.format(prog_name)
    parser = argparse.ArgumentParser(prog=prog_name, description=mesg,
                                     formatter_class=argparse.
                                     RawDescriptionHelpFormatter)
    mut1 = parser.add_mutually_exclusive_group()
    parser.add_argument('-v', '--version', action='version',
                        version='pakit {0}\nALPHA!'.format(__version__))
    parser.add_argument('-a', '--available', default=False,
                        action='store_true', help='list available recipes')
    parser.add_argument('--available-short', default=False,
                        action='store_true',
                        help='list available recipes, terse output')
    parser.add_argument('-c', '--conf', help='yaml config file')
    parser.add_argument('--create-conf', default=False, action='store_true',
                        help='write the default config to CONF')
    parser.add_argument('-d', '--display', nargs='+', metavar='PROG',
                        help='show detailed information on recipe')
    mut1.add_argument('-i', '--install', nargs='+',
                      metavar='PROG', help='install specified program(s)')
    mut1.add_argument('-k', '--search', nargs='+', metavar='WORD',
                      help='search names & descriptions for WORD')
    parser.add_argument('-l', '--list', default=False,
                        action='store_true', help='list installed programs')
    parser.add_argument('--list-short', default=False, action='store_true',
                        help='list installed recipes, terse output')
    mut1.add_argument('-r', '--remove', nargs='+',
                      metavar='PROG', help='remove specified program(s)')
    mut1.add_argument('-u', '--update', default=False, action='store_true',
                      help='update installed programs')

    return parser

def global_init(config_file):
    """
    Performs global configuration of pakit.

    Must be called before using pakit. Will ...
        - Read user configuration.
        - Initialize the logging system.
        - Populate the recipe database.
        - Create configured folders.
        - Setup pakit man page.

    Args:
        config_file: The YAML configuration filename.

    Returns:
        The loaded config object.
    """
    config = Config(config_file)
    pakit.conf.CONFIG = config
    log_init(config)
    logging.debug('Global Config: %s', config)

    for path in config.get('pakit.paths').values():
        try:
            os.makedirs(path)
        except OSError:
            pass

    prefix = config.get('pakit.paths.prefix')
    pakit.conf.IDB = InstallDB(os.path.join(prefix, 'installed.yaml'))
    logging.debug('InstallDB: %s', pakit.conf.IDB)
    pakit.shell.TMP_DIR = os.path.dirname(prefix)

    recipes = RecipeDB(config)
    default = os.path.join(os.path.dirname(os.path.dirname(__file__)),
                           'pakit_recipes')
    recipes.index(default)
    try:
        recipes.index(config.get('pakit.paths.recipes'))
    except KeyError:
        pass

    link_man_page(config.get('pakit.paths.link'))

    return config

def link_man_page(link_dir):
    """
    Silently links man page into link dir.
    """
    src = os.path.join(os.path.dirname(__file__), 'extra', 'pakit.1')
    dst = os.path.join(link_dir, 'share', 'man', 'man1', 'pakit.1')
    try:
        os.makedirs(os.path.dirname(dst))
    except OSError:
        pass
    try:
        os.symlink(src, dst)
    except OSError:
        pass

def log_init(config):
    """
    Setup project wide logging.

    Specifically this will setup both file and console logging.
    """
    log_file = config.get('pakit.log.file')
    try:
        os.makedirs(os.path.dirname(log_file))
    except OSError:
        pass

    root = logging.getLogger()
    root.setLevel(logging.DEBUG)
    log_fmt = '%(levelname)s %(asctime)s %(threadName)s ' \
              '%(filename)s %(message)s'
    my_fmt = logging.Formatter(fmt=log_fmt, datefmt='[%d/%m %H%M.%S]')

    while len(root.handlers):
        root.removeHandler(root.handlers[0])

    rot = logging.handlers.RotatingFileHandler(log_file,
                                               mode='a',
                                               maxBytes=2 ** 22,
                                               backupCount=4)
    rot.setLevel(logging.DEBUG)
    rot.setFormatter(my_fmt)
    root.addHandler(rot)

    stream = logging.StreamHandler()
    stream.setLevel(logging.ERROR)
    stream.setFormatter(my_fmt)
    root.addHandler(stream)

    # Logger for informing user
    pak = logging.getLogger('pakit')
    while len(pak.handlers):
        pak.removeHandler(pak.handlers[0])

    pak.setLevel(logging.INFO)
    pak_fmt = 'pakit: %(asctime)s %(message)s'
    pak_info = logging.Formatter(fmt=pak_fmt, datefmt='[%H:%M:%S]')
    pak_stream = logging.StreamHandler()
    pak_stream.setFormatter(pak_info)
    pak.addHandler(pak_stream)

def parse_tasks(args):
    """
    Parse the program arguments into a list of Tasks to execute

    Args:
        args: An argparse object to map onto tasks.

    Returns:
        A list of Tasks.
    """
    tasks = []

    if args.install:
        tasks.extend([InstallTask(prog) for prog in args.install])
    if args.remove:
        tasks.extend([RemoveTask(prog) for prog in args.remove])
    if args.update:
        tasks.extend([UpdateTask(prog) for prog, _ in pakit.conf.IDB])
    if args.available:
        tasks.append(ListAvailable(False))
    if args.available_short:
        tasks.append(ListAvailable(True))
    if args.display:
        tasks.extend([DisplayTask(prog) for prog in args.display])
    if args.search:
        tasks.append(SearchTask(RecipeDB().names(desc=True), args.search))
    if args.list:
        tasks.append(ListInstalled(False))
    if args.list_short:
        tasks.append(ListInstalled(True))

    return tasks

def search_for_config(default_config=None):
    """
    Search for the most relevant configuration file.

    Will search from current dir upwards until root for config files matching:
        - .pakit.yml
        - .pakit.yaml
        - pakit.yml
        - pakit.yaml

    If nothing found, search in $HOME and then $HOME/.pakit for same files.

    If still nothing found, return default_config.

    Args:
        default_config: The default config if nothing else present.

    Returns:
        If a path existed, return found filename.
        If no paths matched an existing file return the default_config.
    """
    folders = [os.getcwd()]
    cur_dir = os.path.dirname(os.getcwd())
    while cur_dir != '/':
        folders.append(cur_dir)
        cur_dir = os.path.dirname(cur_dir)
    folders.append(os.path.expanduser('~'))
    folders.append(os.path.expanduser('~/.pakit'))

    for folder in folders:
        for conf in ['.pakit.yml', '.pakit.yaml', 'pakit.yml', 'pakit.yaml']:
            config_file = os.path.join(folder, conf)
            if os.path.exists(config_file):
                return config_file

    return default_config

def write_config(config_file):
    """
    Writes the DEFAULT config to the config file.
    Overwrites the file if present.

    Raises:
        PakitError: File exists and is a directory.
        PakitError: File could not be written to.
    """
    config = Config(config_file)
    try:
        os.remove(config.filename)
    except OSError:
        if os.path.isdir(config.filename):
            raise PakitError('Config path is a directory.')
    try:
        config.reset()
        config.write()
        print('pakit: Wrote default config to: ' + config.filename)
    except (IOError, OSError):
        raise PakitError('Failed to write to ' + config.filename)
    finally:
        sys.exit(0)

# TODO: Path modification during operation by os.environ
def main(argv=None):
    """
    The main entry point for this program.

    Args:
        argv: A list of program options, if None use sys.argv.
    """
    if argv is None:
        argv = sys.argv

    parser = args_parser()
    if len(argv) == 1:
        logging.error('No arguments. What should I do?')
        parser.print_usage()
        sys.exit(1)

    args = parser.parse_args(argv[1:])
    if not args.conf:
        args.conf = search_for_config(os.path.expanduser('~/.pakit.yml'))
    if args.create_conf:
        write_config(args.conf)

    config = global_init(args.conf)
    logging.debug('CLI: %s', args)
    PLOG.info('Loading config from: ' + config.filename)

    try:
        tasks = parse_tasks(args)
        for task in tasks:
            PLOG.info('Running: %s', str(task))
            task.run()

        if len(tasks) == 0 and args.update:
            PLOG.info('Nothing to update.')

    except PakitDBError as exc:
        PLOG.info(str(exc))
    except PakitError as exc:
        logging.error(exc)
        logging.error(traceback.format_exc())
        raise

if __name__ == '__main__':
    main()  # pragma: no cover
<|endoftext|>