<commit_before># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

import json

import numpy as np

class DataModel(object):
    def __init__(self, data):
        self.__dict__.update(self.defaults())
        self.__dict__.update(data)

    def defaults(self):
        return {
            'AnimalFlag': 0,
            'AvGRStreamFC': 0,
            'AvGRStreamN': 0,
            'AvGRStreamP': 0,
            'AvTileDrain': np.zeros(100),
            'RurAreaTotal': 0,
            'UrbAreaTotal': 0,
            'Month': np.zeros(12),
            'Landuse': np.zeros(16),
            'Imper': np.zeros(16),
            'LoadRateImp': np.zeros((16, 3)),
            'LoadRatePerv': np.zeros((16, 3)),
            'DisFract': np.zeros((16, 3)),
            'UrbBMPRed': np.zeros((16, 3)),
            'TotSusSolids': np.zeros((16)),
            'Contaminant': np.zeros(3),
            'DayHrs': np.zeros(12),
            'Grow': np.zeros(12),
            'KV': np.zeros(12),
            'd': np.zeros(12),
            'KVD': np.zeros(12),
            'CV': np.zeros(12),
            'Acoef': np.zeros(12),
            'Area': np.zeros(16),
            'AreaE': np.zeros(16),
            'KLSCP': np.zeros(16),
            'CN': np.zeros(16),
            'KF': np.zeros(16),
            'LS': np.zeros(16),
            'C': np.zeros(16),
            'P': np.zeros(16),
            'NewCN': np.zeros((12, 16)),
            'AntMoist': np.zeros(5),
            'PcntET': np.zeros(12),
            'StreetSweepNo': np.zeros(12),
            'ISRR': np.zeros(6),
            'ISRA': np.zeros(6),
            'NitrConc': np.zeros(16),
            'PhosConc': np.zeros(16),
            'ManNitr': np.zeros(16),
            'ManPhos': np.zeros(16),
            'UrbanNitr': np.zeros(16),
            'UrbanPhos': np.zeros(16),
            'PointNitr': np.zeros(12),
            'PointPhos': np.zeros(12),
            'PointFlow': np.zeros(12),
            'AvStreamFlow': np.zeros(12),
            'AvPrecipitation': np.zeros(12),
            'AvEvapoTrans': np.zeros(12),
            'AvGroundWater': np.zeros(12),
            'AvRunoff': np.zeros(12),
            'AvErosion': np.zeros(12),
            'AvSedYield': np.zeros(12),
            'AFON': np.zeros(12),
            'AFOP': np.zeros(12),
            'AvLoad': np.zeros((12, 3)),
            'AvLuLoad': np.zeros((16, 3)),
            'AvDisLoad': np.zeros((16, 3)),
            'AvLuDisLoad': np.zeros((16, 3)),
            'UrbSedLoad': np.zeros((16, 12)),
            'AvGroundNitr': np.zeros(12),
            'AvGroundPhos': np.zeros(12),
            'AvDisNitr': np.zeros(12),
            'AvTotNitr': np.zeros(12),
            'AvDisPhos': np.zeros(12),
            'AvTotPhos': np.zeros(12),
            'AvLuRunoff': np.zeros(16),
            'AvLuErosion': np.zeros(16),
            'AvLuSedYield': np.zeros(16),
            'AvLuDisNitr': np.zeros(16),
            'AvLuTotNitr': np.zeros(16),
            'AvLuDisPhos': np.zeros(16),
            'AvLuTotPhos': np.zeros(16),
            'BSed': np.zeros(16),
            'UrbanSed': np.zeros(16),
            'UrbanErosion': np.zeros(16),
            'ErosWashoff': np.zeros((16, 12)),
            'QRunoff': np.zeros((16, 12)),
            'AgQRunoff': np.zeros((16, 12)),
            'RurQRunoff': np.zeros((16, 12)),
            'UrbQRunoff': np.zeros((16, 12)),
            'NumPondSys': np.zeros(12),
            'NumNormalSys': np.zeros(12),
            'NumShortSys': np.zeros(12),
            'NumDischargeSys': np.zeros(12),
            'NumSewerSys': np.zeros(12),
            'DailyLoad': np.zeros((50, 12, 31)),
            'SepticsDay': np.zeros(12),
            'MonthlyLoad': np.zeros((12, 31)),

            # Declare the daily values as ReDimensional arrays in
            # to Pesticide components
            'DayPondNitr': np.zeros((12, 31)),
            'DayPondPhos': np.zeros((12, 31)),
            'DayNormNitr': np.zeros((12, 31)),
            'DayNormPhos': np.zeros((12, 31)),
            'WashImperv': np.zeros(16),
            'NetSolidLoad': np.zeros(3),
            'DayShortNitr': np.zeros((12, 31)),
            'DayShortPhos': np.zeros((12, 31)),
            'DayDischargeNitr': np.zeros((12, 31)),
            'DayDischargePhos': np.zeros((12, 31)),
            'DaysYear': np.zeros(40),
            'PestAppMonth1': np.zeros(16),
            'PestAppYear1': np.zeros(16),
            'PestAppDate1': np.zeros(16),
            'PestAppMonth2': np.zeros(16),
            'PestAppYear2': np.zeros(16),
            'PestAppDate2': np.zeros(16),
            'PestShedName': np.zeros(12),
            'PestCropArea': np.zeros(12),
            'PestSoilBd': np.zeros(12),
            'PestSoilAwc': np.zeros(12),
            'PestSoilOm': np.zeros(12),
            'PestCropName': np.zeros(12),
            'PestName1': np.zeros(16),
            'PestRate1': np.zeros(31),
            'PestParamCarbon1': np.zeros(16),
            'PestParamWater1': np.zeros(16),
            'PestDecay1': np.zeros(16),
            'PestHalfLife1': np.zeros(16),
            'PestName2': np.zeros(16),
            'PestRate2': np.zeros(31),
            'PestParamCarbon2': np.zeros(16),
            'PestParamWater2': np.zeros(16),
            'PestDecay2': np.zeros(16),
            'PestHalfLife2': np.zeros(16),
            'AvStreamBankEros': np.zeros(12),
            'AvStreamBankN': np.zeros(12),
            'AvStreamBankP': np.zeros(12),
            'CropPercent': np.zeros(12),
            'PestSoilAwcCm': np.zeros(12),

            # Tile Drainage and Flow Variables
            'vTileDrain': np.zeros(12),
            'StreamWithdrawal': np.zeros(12),
            'GroundWithdrawal': np.zeros(12),
            'AvWithdrawal': np.zeros(12),
            'AvTileDrainN': np.zeros(12),
            'AvTileDrainP': np.zeros(12),
            'AvTileDrainSed': np.zeros(12),
            'AvPtSrcFlow': np.zeros(12),

            # DIMENSION LOCAL SEPTIC SYSTEM MODEL ARRAYS
            'MonthPondNitr': np.zeros(12),
            'MonthPondPhos': np.zeros(12),
            'MonthNormNitr': np.zeros(12),
            'MonthShortNitr': np.zeros(12),
            'MonthShortPhos': np.zeros(12),
            'MonthDischargeNitr': np.zeros(12),
            'MonthDischargePhos': np.zeros(12),

            # ANIMAL FEEDING OPERATIONS VARIABLES
            'AnimalName': np.zeros(9),
            'NumAnimals': np.zeros(9),
            'GrazingAnimal': np.zeros(9),
            'AnimalDailyN': np.zeros(9),
            'AnimalDailyP': np.zeros(9),
            'FCOrgsPerDay': np.zeros(9),
            'AvgAnimalWt': np.zeros(9),

            'Month': np.zeros(12),
            'NGPctManApp': np.zeros(12),
            'NGAppNRate': np.zeros(12),
            'NGAppPRate': np.zeros(12),
            'NGAppFCRate': np.zeros(12),
            'NGPctSoilIncRate': np.zeros(12),
            'NGBarnNRate': np.zeros(12),
            'NGBarnPRate': np.zeros(12),
            'NGBarnFCRate': np.zeros(12),

            'PctGrazing': np.zeros(12),
            'PctStreams': np.zeros(12),
            'GrazingNRate': np.zeros(12),
            'GrazingPRate': np.zeros(12),
            'GrazingFCRate': np.zeros(12),
            'GRPctManApp': np.zeros(12),
            'GRAppNRate': np.zeros(12),
            'GRAppPRate': np.zeros(12),
            'GRAppFCRate': np.zeros(12),
            'GRPctSoilIncRate': np.zeros(12),
            'GRBarnNRate': np.zeros(12),
            'GRBarnPRate': np.zeros(12),
            'GRBarnFCRate': np.zeros(12),

            # Calculated Values for Animal Feeding Operations
            'NGLoadN': np.zeros(9),
            'NGLoadP': np.zeros(9),
            'NGLoadFC': np.zeros(9),
            'NGAccManAppN': np.zeros(12),
            'NGAccManAppP': np.zeros(12),
            'NGAccManAppFC': np.zeros(12),
            'NGAppManN': np.zeros(12),
            'NGInitBarnN': np.zeros(12),
            'NGAppManP': np.zeros(12),
            'NGInitBarnP': np.zeros(12),
            'NGAppManFC': np.zeros(12),
            'NGInitBarnFC': np.zeros(12),

            'GRLoadN': np.zeros(9),
            'GRLoadP': np.zeros(9),
            'GRLoadFC': np.zeros(9),
            'GRAccManAppN': np.zeros(12),
            'GRAccManAppP': np.zeros(12),
            'GRAccManAppFC': np.zeros(12),
            'GRAppManN': np.zeros(12),
            'GRInitBarnN': np.zeros(12),
            'GRAppManP': np.zeros(12),
            'GRInitBarnP': np.zeros(12),
            'GRAppManFC': np.zeros(12),
            'GRInitBarnFC': np.zeros(12),
            'GrazingN': np.zeros(12),
            'GrazingP': np.zeros(12),
            'GrazingFC': np.zeros(12),
            'GRStreamN': np.zeros(12),
            'GRStreamP': np.zeros(12),
            'GRStreamFC': np.zeros(12),
            'AvAnimalN': np.zeros(12),
            'AvAnimalP': np.zeros(12),
            'AvAnimalFC': np.zeros(12),
            'AvWWOrgs': np.zeros(12),
            'AvSSOrgs': np.zeros(12),
            'AvUrbOrgs': np.zeros(12),
            'AvWildOrgs': np.zeros(12),
            'AvTotalOrgs': np.zeros(12),
            'AvCMStream': np.zeros(12),
            'AvOrgConc': np.zeros(12),
            'AvGRLostBarnN': np.zeros(12),
            'AvGRLostBarnP': np.zeros(12),
            'AvNGLostBarnN': np.zeros(12),
            'AvNGLostBarnP': np.zeros(12),
            'AvNGLostManP': np.zeros(12),
            'AvNGLostBarnFC': np.zeros(12),
            'AvGRLostBarnFC': np.zeros(12),

            # Initialize non-arrayed variables
            'Nqual': 0,
            'q': 0,
            'k': 0,
            'BasinArea': 0,
            'Qretention': 0,
            'FilterEff': 0,
            'FilterWidth': 0,
            'OutFiltWidth': 0,
            'Capacity': 0,
            'BasinDeadStorage': 0,
            'DaysToDrain': 0,
            'CleanMon': 0,
            'Clean': 0,
            'CleanSwitch': 0,
            'OutletCoef': 0,
            'BasinVol': 0,
            'Volume': 0,
            'ActiveVol': 0,
            'DetentFlow': 0,
            'AnnDayHrs': 0,
            'AreaTotal': 0,
            'SweepFrac': np.zeros(12),

            # CNI and CNP are 1-based in VB, so although the first dimension
            # should be size 3, we increase it by 1 because the model
            # references indexes 1, 2, and 3. The first index is never used.
            'CNI': np.zeros((3 + 1, 16)),
            'CNP': np.zeros((3 + 1, 16)),

            # Load Delivery data
            'AttenFlowDist': 0,
            'AttenFlowVel': 0,
            'AttenLossRateN': 0,
            'AttenLossRateP': 0,
            'AttenLossRateTSS': 0,
            'AttenLossRatePath': 0,
            'StreamFlowVolAdj': 1,
        }

    def __str__(self):
        return '<GWLF-E DataModel>'

    def tojson(self):
        return json.dumps(self.__dict__, cls=NumpyAwareJSONEncoder)

class NumpyAwareJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return json.JSONEncoder.default(self, obj)
<commit_msg>Add support for serializing Decimals<commit_after># -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

from decimal import Decimal
import json

import numpy as np

class DataModel(object):
    def __init__(self, data):
        self.__dict__.update(self.defaults())
        self.__dict__.update(data)

    def defaults(self):
        return {
            'AnimalFlag': 0,
            'AvGRStreamFC': 0,
            'AvGRStreamN': 0,
            'AvGRStreamP': 0,
            'AvTileDrain': np.zeros(100),
            'RurAreaTotal': 0,
            'UrbAreaTotal': 0,
            'Month': np.zeros(12),
            'Landuse': np.zeros(16),
            'Imper': np.zeros(16),
            'LoadRateImp': np.zeros((16, 3)),
            'LoadRatePerv': np.zeros((16, 3)),
            'DisFract': np.zeros((16, 3)),
            'UrbBMPRed': np.zeros((16, 3)),
            'TotSusSolids': np.zeros((16)),
            'Contaminant': np.zeros(3),
            'DayHrs': np.zeros(12),
            'Grow': np.zeros(12),
            'KV': np.zeros(12),
            'd': np.zeros(12),
            'KVD': np.zeros(12),
            'CV': np.zeros(12),
            'Acoef': np.zeros(12),
            'Area': np.zeros(16),
            'AreaE': np.zeros(16),
            'KLSCP': np.zeros(16),
            'CN': np.zeros(16),
            'KF': np.zeros(16),
            'LS': np.zeros(16),
            'C': np.zeros(16),
            'P': np.zeros(16),
            'NewCN': np.zeros((12, 16)),
            'AntMoist': np.zeros(5),
            'PcntET': np.zeros(12),
            'StreetSweepNo': np.zeros(12),
            'ISRR': np.zeros(6),
            'ISRA': np.zeros(6),
            'NitrConc': np.zeros(16),
            'PhosConc': np.zeros(16),
            'ManNitr': np.zeros(16),
            'ManPhos': np.zeros(16),
            'UrbanNitr': np.zeros(16),
            'UrbanPhos': np.zeros(16),
            'PointNitr': np.zeros(12),
            'PointPhos': np.zeros(12),
            'PointFlow': np.zeros(12),
            'AvStreamFlow': np.zeros(12),
            'AvPrecipitation': np.zeros(12),
            'AvEvapoTrans': np.zeros(12),
            'AvGroundWater': np.zeros(12),
            'AvRunoff': np.zeros(12),
            'AvErosion': np.zeros(12),
            'AvSedYield': np.zeros(12),
            'AFON': np.zeros(12),
            'AFOP': np.zeros(12),
            'AvLoad': np.zeros((12, 3)),
            'AvLuLoad': np.zeros((16, 3)),
            'AvDisLoad': np.zeros((16, 3)),
            'AvLuDisLoad': np.zeros((16, 3)),
            'UrbSedLoad': np.zeros((16, 12)),
            'AvGroundNitr': np.zeros(12),
            'AvGroundPhos': np.zeros(12),
            'AvDisNitr': np.zeros(12),
            'AvTotNitr': np.zeros(12),
            'AvDisPhos': np.zeros(12),
            'AvTotPhos': np.zeros(12),
            'AvLuRunoff': np.zeros(16),
            'AvLuErosion': np.zeros(16),
            'AvLuSedYield': np.zeros(16),
            'AvLuDisNitr': np.zeros(16),
            'AvLuTotNitr': np.zeros(16),
            'AvLuDisPhos': np.zeros(16),
            'AvLuTotPhos': np.zeros(16),
            'BSed': np.zeros(16),
            'UrbanSed': np.zeros(16),
            'UrbanErosion': np.zeros(16),
            'ErosWashoff': np.zeros((16, 12)),
            'QRunoff': np.zeros((16, 12)),
            'AgQRunoff': np.zeros((16, 12)),
            'RurQRunoff': np.zeros((16, 12)),
            'UrbQRunoff': np.zeros((16, 12)),
            'NumPondSys': np.zeros(12),
            'NumNormalSys': np.zeros(12),
            'NumShortSys': np.zeros(12),
            'NumDischargeSys': np.zeros(12),
            'NumSewerSys': np.zeros(12),
            'DailyLoad': np.zeros((50, 12, 31)),
            'SepticsDay': np.zeros(12),
            'MonthlyLoad': np.zeros((12, 31)),

            # Declare the daily values as ReDimensional arrays in
            # to Pesticide components
            'DayPondNitr': np.zeros((12, 31)),
            'DayPondPhos': np.zeros((12, 31)),
            'DayNormNitr': np.zeros((12, 31)),
            'DayNormPhos': np.zeros((12, 31)),
            'WashImperv': np.zeros(16),
            'NetSolidLoad': np.zeros(3),
            'DayShortNitr': np.zeros((12, 31)),
            'DayShortPhos': np.zeros((12, 31)),
            'DayDischargeNitr': np.zeros((12, 31)),
            'DayDischargePhos': np.zeros((12, 31)),
            'DaysYear': np.zeros(40),
            'PestAppMonth1': np.zeros(16),
            'PestAppYear1': np.zeros(16),
            'PestAppDate1': np.zeros(16),
            'PestAppMonth2': np.zeros(16),
            'PestAppYear2': np.zeros(16),
            'PestAppDate2': np.zeros(16),
            'PestShedName': np.zeros(12),
            'PestCropArea': np.zeros(12),
            'PestSoilBd': np.zeros(12),
            'PestSoilAwc': np.zeros(12),
            'PestSoilOm': np.zeros(12),
            'PestCropName': np.zeros(12),
            'PestName1': np.zeros(16),
            'PestRate1': np.zeros(31),
            'PestParamCarbon1': np.zeros(16),
            'PestParamWater1': np.zeros(16),
            'PestDecay1': np.zeros(16),
            'PestHalfLife1': np.zeros(16),
            'PestName2': np.zeros(16),
            'PestRate2': np.zeros(31),
            'PestParamCarbon2': np.zeros(16),
            'PestParamWater2': np.zeros(16),
            'PestDecay2': np.zeros(16),
            'PestHalfLife2': np.zeros(16),
            'AvStreamBankEros': np.zeros(12),
            'AvStreamBankN': np.zeros(12),
            'AvStreamBankP': np.zeros(12),
            'CropPercent': np.zeros(12),
            'PestSoilAwcCm': np.zeros(12),

            # Tile Drainage and Flow Variables
            'vTileDrain': np.zeros(12),
            'StreamWithdrawal': np.zeros(12),
            'GroundWithdrawal': np.zeros(12),
            'AvWithdrawal': np.zeros(12),
            'AvTileDrainN': np.zeros(12),
            'AvTileDrainP': np.zeros(12),
            'AvTileDrainSed': np.zeros(12),
            'AvPtSrcFlow': np.zeros(12),

            # DIMENSION LOCAL SEPTIC SYSTEM MODEL ARRAYS
            'MonthPondNitr': np.zeros(12),
            'MonthPondPhos': np.zeros(12),
            'MonthNormNitr': np.zeros(12),
            'MonthShortNitr': np.zeros(12),
            'MonthShortPhos': np.zeros(12),
            'MonthDischargeNitr': np.zeros(12),
            'MonthDischargePhos': np.zeros(12),

            # ANIMAL FEEDING OPERATIONS VARIABLES
            'AnimalName': np.zeros(9),
            'NumAnimals': np.zeros(9),
            'GrazingAnimal': np.zeros(9),
            'AnimalDailyN': np.zeros(9),
            'AnimalDailyP': np.zeros(9),
            'FCOrgsPerDay': np.zeros(9),
            'AvgAnimalWt': np.zeros(9),

            'Month': np.zeros(12),
            'NGPctManApp': np.zeros(12),
            'NGAppNRate': np.zeros(12),
            'NGAppPRate': np.zeros(12),
            'NGAppFCRate': np.zeros(12),
            'NGPctSoilIncRate': np.zeros(12),
            'NGBarnNRate': np.zeros(12),
            'NGBarnPRate': np.zeros(12),
            'NGBarnFCRate': np.zeros(12),

            'PctGrazing': np.zeros(12),
            'PctStreams': np.zeros(12),
            'GrazingNRate': np.zeros(12),
            'GrazingPRate': np.zeros(12),
            'GrazingFCRate': np.zeros(12),
            'GRPctManApp': np.zeros(12),
            'GRAppNRate': np.zeros(12),
            'GRAppPRate': np.zeros(12),
            'GRAppFCRate': np.zeros(12),
            'GRPctSoilIncRate': np.zeros(12),
            'GRBarnNRate': np.zeros(12),
            'GRBarnPRate': np.zeros(12),
            'GRBarnFCRate': np.zeros(12),

            # Calculated Values for Animal Feeding Operations
            'NGLoadN': np.zeros(9),
            'NGLoadP': np.zeros(9),
            'NGLoadFC': np.zeros(9),
            'NGAccManAppN': np.zeros(12),
            'NGAccManAppP': np.zeros(12),
            'NGAccManAppFC': np.zeros(12),
            'NGAppManN': np.zeros(12),
            'NGInitBarnN': np.zeros(12),
            'NGAppManP': np.zeros(12),
            'NGInitBarnP': np.zeros(12),
            'NGAppManFC': np.zeros(12),
            'NGInitBarnFC': np.zeros(12),

            'GRLoadN': np.zeros(9),
            'GRLoadP': np.zeros(9),
            'GRLoadFC': np.zeros(9),
            'GRAccManAppN': np.zeros(12),
            'GRAccManAppP': np.zeros(12),
            'GRAccManAppFC': np.zeros(12),
            'GRAppManN': np.zeros(12),
            'GRInitBarnN': np.zeros(12),
            'GRAppManP': np.zeros(12),
            'GRInitBarnP': np.zeros(12),
            'GRAppManFC': np.zeros(12),
            'GRInitBarnFC': np.zeros(12),
            'GrazingN': np.zeros(12),
            'GrazingP': np.zeros(12),
            'GrazingFC': np.zeros(12),
            'GRStreamN': np.zeros(12),
            'GRStreamP': np.zeros(12),
            'GRStreamFC': np.zeros(12),
            'AvAnimalN': np.zeros(12),
            'AvAnimalP': np.zeros(12),
            'AvAnimalFC': np.zeros(12),
            'AvWWOrgs': np.zeros(12),
            'AvSSOrgs': np.zeros(12),
            'AvUrbOrgs': np.zeros(12),
            'AvWildOrgs': np.zeros(12),
            'AvTotalOrgs': np.zeros(12),
            'AvCMStream': np.zeros(12),
            'AvOrgConc': np.zeros(12),
            'AvGRLostBarnN': np.zeros(12),
            'AvGRLostBarnP': np.zeros(12),
            'AvNGLostBarnN': np.zeros(12),
            'AvNGLostBarnP': np.zeros(12),
            'AvNGLostManP': np.zeros(12),
            'AvNGLostBarnFC': np.zeros(12),
            'AvGRLostBarnFC': np.zeros(12),

            # Initialize non-arrayed variables
            'Nqual': 0,
            'q': 0,
            'k': 0,
            'BasinArea': 0,
            'Qretention': 0,
            'FilterEff': 0,
            'FilterWidth': 0,
            'OutFiltWidth': 0,
            'Capacity': 0,
            'BasinDeadStorage': 0,
            'DaysToDrain': 0,
            'CleanMon': 0,
            'Clean': 0,
            'CleanSwitch': 0,
            'OutletCoef': 0,
            'BasinVol': 0,
            'Volume': 0,
            'ActiveVol': 0,
            'DetentFlow': 0,
            'AnnDayHrs': 0,
            'AreaTotal': 0,
            'SweepFrac': np.zeros(12),

            # CNI and CNP are 1-based in VB, so although the first dimension
            # should be size 3, we increase it by 1 because the model
            # references indexes 1, 2, and 3. The first index is never used.
            'CNI': np.zeros((3 + 1, 16)),
            'CNP': np.zeros((3 + 1, 16)),

            # Load Delivery data
            'AttenFlowDist': 0,
            'AttenFlowVel': 0,
            'AttenLossRateN': 0,
            'AttenLossRateP': 0,
            'AttenLossRateTSS': 0,
            'AttenLossRatePath': 0,
            'StreamFlowVolAdj': 1,
        }

    def __str__(self):
        return '<GWLF-E DataModel>'

    def tojson(self):
        return json.dumps(self.__dict__, cls=NumpyAwareJSONEncoder)

class NumpyAwareJSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, Decimal):
            return float(obj)
        return json.JSONEncoder.default(self, obj)
<|endoftext|>


<commit_before># coding: utf8

from __future__ import unicode_literals
import codecs
from nltk.parse import DependencyGraph

def coarse_pos_e(tags):
	"""
	Coarse POS tags of Dadegan corpus:
		N: Noun, V: Verb, ADJ: Adjective, ADV: Adverb, PR: Pronoun, PREP: Preposition, POSTP: Postposition, CONJ: Conjunction, PUNC: Punctuation, ADR: Address Term, IDEN: Title, PART: Particle, POSNUM: Post-noun Modifier, PREM: Pre-modifier, PRENUM: Pre-noun Numeral, PSUS: Pseudo-sentence, SUBR: Subordinating Clause

	>>> coarse_pos_e(['N', 'IANM'])
	'N'
	"""

	map = {'N': 'N', 'V': 'V', 'ADJ': 'AJ', 'ADV': 'ADV', 'PR': 'PRO', 'PREM': 'DET', 'PREP': 'P', 'POSTP': 'POSTP', 'PRENUM': 'NUM', 'CONJ': 'CONJ', 'PUNC': 'PUNC'}
	return map.get(tags[0], '') + ('e' if 'EZ' in tags else '')

class DadeganReader():
	"""
	interfaces [Persian Dependency Treebank](http://dadegan.ir/perdt/download)
	"""

	def __init__(self, conll_file='corpora/dadegan.conll', pos_map=coarse_pos_e):
		self._conll_file = conll_file
		self._pos_map = pos_map if pos_map else lambda tags: ','.join(tags)

	def _sentences(self):
		text = codecs.open(self._conll_file, encoding='utf8').read()

		# refine text
		text = text.replace('','').replace('\t','\t').replace('\t','\t').replace('\t ','\t').replace(' \t','\t').replace('\r', '').replace('\u2029', '')

		for item in text.replace(' ', '_').split('\n\n'):
			if item.strip():
				yield item

	def trees(self):
		for sentence in self._sentences():
			tree = DependencyGraph(sentence)

			for node in tree.nodelist[1:]:
				node['mtag'] = [node['ctag'], node['tag']]

			for node in tree.nodelist[1:]:
				if node['rel'] in ('MOZ', 'NPOSTMOD'):
					tree.nodelist[node['head']]['mtag'].append('EZ')

			for node in tree.nodelist[1:]:
				node['mtag'] = self._pos_map(node['mtag'])

			yield tree

	def sents(self):
		"""
		>>> next(dadegan.sents())
		[('دنیای', 'Ne'), ('آدولف', 'N'), ('بورن', 'N'), ('دنیای', 'Ne'), ('اتفاقات', 'Ne'), ('رویایی', 'AJ'), ('است', 'V'), ('.', 'PUNC')]
		"""
		for tree in self.trees():
			yield [(node['word'], node['mtag']) for node in tree.nodelist[1:]]

	def chunked_trees(self):
		"""
		>>> tree2brackets(next(dadegan.chunked_trees()))
		'[این میهمانی NP] [به PP] [منظور آشنایی هم‌تیمی‌های او NP] [با PP] [غذاهای ایرانی NP] [ترتیب داده شد VP] .'
		"""
<commit_msg>Add Chunked_sents<commit_after># coding: utf8

from __future__ import unicode_literals
import codecs
from nltk.parse import DependencyGraph
from nltk.tree import Tree
from .Chunker import tree2brackets

def coarse_pos_e(tags):
	"""
	Coarse POS tags of Dadegan corpus:
		N: Noun, V: Verb, ADJ: Adjective, ADV: Adverb, PR: Pronoun, PREP: Preposition, POSTP: Postposition, CONJ: Conjunction, PUNC: Punctuation, ADR: Address Term, IDEN: Title, PART: Particle, POSNUM: Post-noun Modifier, PREM: Pre-modifier, PRENUM: Pre-noun Numeral, PSUS: Pseudo-sentence, SUBR: Subordinating Clause

	>>> coarse_pos_e(['N', 'IANM'])
	'N'
	"""

	map = {'N': 'N', 'V': 'V', 'ADJ': 'AJ', 'ADV': 'ADV', 'PR': 'PRO', 'PREM': 'DET', 'PREP': 'P', 'POSTP': 'POSTP', 'PRENUM': 'NUM', 'CONJ': 'CONJ', 'PUNC': 'PUNC'}
	return map.get(tags[0], '') + ('e' if 'EZ' in tags else '')

class DadeganReader():
	"""
	interfaces [Persian Dependency Treebank](http://dadegan.ir/perdt/download)
	"""

	def __init__(self, conll_file='corpora/dadegan.conll', pos_map=coarse_pos_e):
		self._conll_file = conll_file
		self._pos_map = pos_map if pos_map else lambda tags: ','.join(tags)

	def _sentences(self):
		text = codecs.open(self._conll_file, encoding='utf8').read()

		# refine text
		text = text.replace('','').replace('\t','\t').replace('\t','\t').replace('\t ','\t').replace(' \t','\t').replace('\r', '').replace('\u2029', '')

		for item in text.replace(' ', '_').split('\n\n'):
			if item.strip():
				yield item

	def trees(self):
		for sentence in self._sentences():
			tree = DependencyGraph(sentence)

			for node in tree.nodelist[1:]:
				node['mtag'] = [node['ctag'], node['tag']]

			for node in tree.nodelist[1:]:
				if node['rel'] in ('MOZ', 'NPOSTMOD'):
					tree.nodelist[node['head']]['mtag'].append('EZ')

			for node in tree.nodelist[1:]:
				node['mtag'] = self._pos_map(node['mtag'])

			yield tree

	def sents(self):
		"""
		>>> next(dadegan.sents())
		[('این', 'DET'), ('میهمانی', 'N'), ('به', 'P'), ('منظور', 'Ne'), ('آشنایی', 'Ne'), ('هم‌تیمی‌های', 'Ne'), ('او', 'PRO'), ('با', 'P'), ('غذاهای', 'Ne'), ('ایرانی', 'AJ'), ('ترتیب', 'N'), ('داده_شد', 'V'), ('.', 'PUNC')]
		"""
		for tree in self.trees():
			yield [(node['word'], node['mtag']) for node in tree.nodelist[1:]]

	def chunked_trees(self):
		"""
		>>> tree2brackets(next(dadegan.chunked_trees()))
		'[این میهمانی NP] [به PP] [منظور آشنایی هم‌تیمی‌های او NP] [با PP] [غذاهای ایرانی NP] [ترتیب داده_شد VP] .'
		"""
		trees = self.trees()
		for tree in trees:
			chunks = []
			nodelist = tree.nodelist
			for i in range(1, len(nodelist)):
				node = (nodelist[i]['word'], nodelist[i]['mtag'])
				appended = False
				if nodelist[i]['ctag'] in ['PREP']:
					for j in nodelist[i]['deps']:
						if j == i - 1 and type(chunks[-1]) == Tree and chunks[-1].label() == 'PP':
							chunks[-1].append(node)
							appended = True
					if nodelist[i]['head'] == i - 1 and chunks[-1].label() != Tree and chunks[-1].label() == 'PP':
						chunks[-1].append(node)
						appended = True
					if not appended:
						chunks.append(Tree('PP', [node]))
				elif nodelist[i]['ctag'] in ['PUNC']:
					chunks.append(node)
				elif nodelist[i]['ctag'] in ['N', 'PREM', 'ADJ', 'PR' ]:
					for j in nodelist[i]['deps']:
						if j == i - 1 and type(chunks[-1]) == Tree and chunks[-1].label() != 'PP':
							chunks[-1].append(node)
							appended = True
					if nodelist[i]['head'] == i - 1 and chunks[-1].label() != Tree and chunks[-1].label() != 'PP':
						chunks[-1].append(node)
						appended = True
					if not appended:
						chunks.append(Tree('NP', [node]))
				elif nodelist[i]['ctag'] in ['V']:
					appended = False
					for j in nodelist[i]['deps']:
						if j == i - 1 and type(chunks[-1]) == Tree:
							leaves = chunks.pop().leaves()
							leaves.append(node)
							chunks.append(Tree('VP', leaves))
							appended = True
					if not appended:
						chunks.append(Tree('VP', [node]))
			yield Tree('S', chunks)

<|endoftext|>


<commit_before>"""
Supplies Layer and related classes that allow overlaying of Views,
including Overlay. A Layer is the final extension of View base class
that allows Views to be overlayed on top of each other.

Also supplies ViewMap which is the primary multi-dimensional Map type
for indexing, slicing and animating collections of Views.
"""

from functools import reduce
import numpy as np

import param
from .dimension import Dimension, Dimensioned, ViewableElement
from .ndmapping import UniformNdMapping
from .layout import Composable, Layout
from .util import sanitize_identifier, unique_array

class Overlayable(object):
    """
    Overlayable provides a mix-in class to support the
    mul operation for overlaying multiple elements.
    """

    def __mul__(self, other):
        if type(other).__name__ == 'DynamicMap':
            from ..util import Dynamic
            def dynamic_mul(element):
                return self * element
            return Dynamic(other, operation=dynamic_mul)
        if isinstance(other, UniformNdMapping) and not isinstance(other, CompositeOverlay):
            items = [(k, self * v) for (k, v) in other.items()]
            return other.clone(items)

        self_item = [((self.group, self.label if self.label else 'I'), self)]
        other_items = (other.items() if isinstance(other, Overlay)
                       else [((other.group, other.label if other.label else 'I'), other)])
        return Overlay(items=Overlay.relabel_item_paths(list(self_item) + list(other_items)))

class CompositeOverlay(ViewableElement, Composable):
    """
    CompositeOverlay provides a common baseclass for Overlay classes.
    """

    _deep_indexable = True

    def hist(self, index=0, adjoin=True, dimension=None, **kwargs):
        valid_ind = isinstance(index, int) and (0 <= index < len(self))
        valid_label = index in [el.label for el in self]
        if not any([valid_ind, valid_label]):
            raise TypeError("Please supply a suitable index or label for the histogram data")

        hists = self.get(index).hist(adjoin=False, dimension=dimension, **kwargs)
        if not isinstance(hists, Layout):
            hists = [hists]
        if not isinstance(dimension, list):
            dimension = ['Default']
        if adjoin:
            layout = self
            for hist in hists:
                layout = layout << hist
            layout.main_layer = index
        elif len(dimension) > 1:
            layout = hists
        else:
            layout = hists[0]
        return layout

    def dimension_values(self, dimension, expanded=True, flat=True):
        values = []
        found = False
        for el in self:
            if dimension in el.dimensions(label=True):
                values.append(el.dimension_values(dimension))
                found = True
        if not found:
            return super(CompositeOverlay, self).dimension_values(dimension, expanded, flat)
        values = [v for v in values if v is not None and len(v)]
        if not values:
            return np.array()
        vals = np.concatenate(values)
        return vals if expanded else unique_array(vals)

class Overlay(Layout, CompositeOverlay):
    """
    An Overlay consists of multiple Views (potentially of
    heterogeneous type) presented one on top each other with a
    particular z-ordering.

    Overlays along with Views constitute the only valid leaf types of
    a Layout and in fact extend the Layout structure. Overlays are
    constructed using the * operator (building an identical structure
    to the + operator) and are the only objects that inherit both from
    Layout and CompositeOverlay.
    """

    @classmethod
    def _from_values(cls, val):
        return reduce(lambda x,y: x*y, val).map(lambda x: x.display('auto'), [Overlay])

    def __init__(self, items=None, group=None, label=None, **params):
        view_params = ViewableElement.params().keys()
        self.__dict__['_fixed'] = False
        self.__dict__['_group'] = group
        self.__dict__['_label'] = label
        Layout.__init__(self, items,
                          **{k:v for k,v in params.items() if k not in view_params})
        ViewableElement.__init__(self, self.data,
                                 **{k:v for k,v in params.items() if k in view_params})

    def __getitem__(self, key):
        """
        Allows transparently slicing the Elements in the Overlay
        to select specific layers in an Overlay use the .get method.
        """
        return Overlay([(k, v[key]) for k, v in self.items()])

    def get(self, identifier, default=None):
        if isinstance(identifier, int):
            values = list(self.data.values())
            if 0 <= identifier < len(values):
                return values[identifier]
            else:
                return default
        return super(Overlay, self).get(identifier, default)

    def __add__(self, other):
        return Layout.from_values(self) + Layout.from_values(other)

    def __mul__(self, other):
        if isinstance(other, Overlay):
            items = list(self.data.items()) + list(other.data.items())
        elif isinstance(other, ViewableElement):
            label = other.label if other.label else 'I'
            items = list(self.data.items()) + [((other.group, label), other)]
        elif isinstance(other, UniformNdMapping):
            raise NotImplementedError

        return Overlay(items=self.relabel_item_paths(items)).display('all')

    def collapse(self, function):
        """
        Collapses all the Elements in the Overlay using the
        supplied function if they share a common type and group.
        """
        elements = list(self)
        types = [type(el) for el in elements]
        values = [el.group for el in elements]
        if not len(set(types)) == 1 and len(set(values)) == 1:
            raise Exception("Overlay is not homogenous in type or group "
                            "and cannot be collapsed.")
        else:
            return elements[0].clone(types[0].collapse_data([el.data for el in elements],
                                                            function, self.kdims))

    @property
    def group(self):
        if self._group:
            return self._group
        elements = [el for el in self if not el._auxiliary_component]
        values = {el.group for el in elements}
        types = {type(el) for el in elements}
        if values:
            group = list(values)[0]
            vtype = list(types)[0].__name__
        else:
            group, vtype = [], ''
        if len(values) == 1 and group != vtype:
            return group
        else:
            return type(self).__name__

    @group.setter
    def group(self, group):
        if not sanitize_identifier.allowable(group):
            raise ValueError("Supplied group %s contains invalid characters." %
                             group)
        else:
            self._group = group

    @property
    def label(self):
        if self._label:
            return self._label
        labels = {el.label for el in self
                  if not el._auxiliary_component}
        if len(labels) == 1:
            return list(labels)[0]
        else:
            return ''

    @label.setter
    def label(self, label):
        if not sanitize_identifier.allowable(label):
            raise ValueError("Supplied group %s contains invalid characters." %
                             label)
        self._label = label

    @property
    def ddims(self):
        dimensions = []
        dimension_names = []
        for el in self:
            for dim in el.dimensions():
                if dim.name not in dimension_names:
                    dimensions.append(dim)
                    dimension_names.append(dim.name)
        return dimensions

    @property
    def shape(self):
        raise NotImplementedError

class NdOverlay(UniformNdMapping, CompositeOverlay, Overlayable):
    """
    An NdOverlay allows a group of NdOverlay to be overlaid together. NdOverlay can
    be indexed out of an overlay and an overlay is an iterable that iterates
    over the contained layers.
    """

    kdims = param.List(default=[Dimension('Element')], constant=True, doc="""
        List of dimensions the NdOverlay can be indexed by.""")

    _deep_indexable = True

    def __init__(self, overlays=None, **params):
        super(NdOverlay, self).__init__(overlays, **params)

__all__ = list(set([_k for _k, _v in locals().items()
                    if isinstance(_v, type) and issubclass(_v, Dimensioned)])) + ['Overlayable']
<commit_msg>Added collate method on Overlay<commit_after>"""
Supplies Layer and related classes that allow overlaying of Views,
including Overlay. A Layer is the final extension of View base class
that allows Views to be overlayed on top of each other.

Also supplies ViewMap which is the primary multi-dimensional Map type
for indexing, slicing and animating collections of Views.
"""

from functools import reduce
import numpy as np

import param
from .dimension import Dimension, Dimensioned, ViewableElement
from .ndmapping import UniformNdMapping
from .layout import Composable, Layout
from .util import sanitize_identifier, unique_array

class Overlayable(object):
    """
    Overlayable provides a mix-in class to support the
    mul operation for overlaying multiple elements.
    """

    def __mul__(self, other):
        if type(other).__name__ == 'DynamicMap':
            from ..util import Dynamic
            def dynamic_mul(element):
                return self * element
            return Dynamic(other, operation=dynamic_mul)
        if isinstance(other, UniformNdMapping) and not isinstance(other, CompositeOverlay):
            items = [(k, self * v) for (k, v) in other.items()]
            return other.clone(items)

        self_item = [((self.group, self.label if self.label else 'I'), self)]
        other_items = (other.items() if isinstance(other, Overlay)
                       else [((other.group, other.label if other.label else 'I'), other)])
        return Overlay(items=Overlay.relabel_item_paths(list(self_item) + list(other_items)))

class CompositeOverlay(ViewableElement, Composable):
    """
    CompositeOverlay provides a common baseclass for Overlay classes.
    """

    _deep_indexable = True

    def hist(self, index=0, adjoin=True, dimension=None, **kwargs):
        valid_ind = isinstance(index, int) and (0 <= index < len(self))
        valid_label = index in [el.label for el in self]
        if not any([valid_ind, valid_label]):
            raise TypeError("Please supply a suitable index or label for the histogram data")

        hists = self.get(index).hist(adjoin=False, dimension=dimension, **kwargs)
        if not isinstance(hists, Layout):
            hists = [hists]
        if not isinstance(dimension, list):
            dimension = ['Default']
        if adjoin:
            layout = self
            for hist in hists:
                layout = layout << hist
            layout.main_layer = index
        elif len(dimension) > 1:
            layout = hists
        else:
            layout = hists[0]
        return layout

    def dimension_values(self, dimension, expanded=True, flat=True):
        values = []
        found = False
        for el in self:
            if dimension in el.dimensions(label=True):
                values.append(el.dimension_values(dimension))
                found = True
        if not found:
            return super(CompositeOverlay, self).dimension_values(dimension, expanded, flat)
        values = [v for v in values if v is not None and len(v)]
        if not values:
            return np.array()
        vals = np.concatenate(values)
        return vals if expanded else unique_array(vals)

class Overlay(Layout, CompositeOverlay):
    """
    An Overlay consists of multiple Views (potentially of
    heterogeneous type) presented one on top each other with a
    particular z-ordering.

    Overlays along with Views constitute the only valid leaf types of
    a Layout and in fact extend the Layout structure. Overlays are
    constructed using the * operator (building an identical structure
    to the + operator) and are the only objects that inherit both from
    Layout and CompositeOverlay.
    """

    @classmethod
    def _from_values(cls, val):
        return reduce(lambda x,y: x*y, val).map(lambda x: x.display('auto'), [Overlay])

    def __init__(self, items=None, group=None, label=None, **params):
        view_params = ViewableElement.params().keys()
        self.__dict__['_fixed'] = False
        self.__dict__['_group'] = group
        self.__dict__['_label'] = label
        Layout.__init__(self, items,
                          **{k:v for k,v in params.items() if k not in view_params})
        ViewableElement.__init__(self, self.data,
                                 **{k:v for k,v in params.items() if k in view_params})

    def __getitem__(self, key):
        """
        Allows transparently slicing the Elements in the Overlay
        to select specific layers in an Overlay use the .get method.
        """
        return Overlay([(k, v[key]) for k, v in self.items()])

    def get(self, identifier, default=None):
        if isinstance(identifier, int):
            values = list(self.data.values())
            if 0 <= identifier < len(values):
                return values[identifier]
            else:
                return default
        return super(Overlay, self).get(identifier, default)

    def __add__(self, other):
        return Layout.from_values(self) + Layout.from_values(other)

    def __mul__(self, other):
        if isinstance(other, Overlay):
            items = list(self.data.items()) + list(other.data.items())
        elif isinstance(other, ViewableElement):
            label = other.label if other.label else 'I'
            items = list(self.data.items()) + [((other.group, label), other)]
        elif isinstance(other, UniformNdMapping):
            raise NotImplementedError

        return Overlay(items=self.relabel_item_paths(items)).display('all')

    def collate(self):
        """
        Collates any objects in the Overlay resolving any issues
        the recommended nesting structure.
        """
        return reduce(lambda x,y: x*y, self.values())

    def collapse(self, function):
        """
        Collapses all the Elements in the Overlay using the
        supplied function if they share a common type and group.
        """
        elements = list(self)
        types = [type(el) for el in elements]
        values = [el.group for el in elements]
        if not len(set(types)) == 1 and len(set(values)) == 1:
            raise Exception("Overlay is not homogenous in type or group "
                            "and cannot be collapsed.")
        else:
            return elements[0].clone(types[0].collapse_data([el.data for el in elements],
                                                            function, self.kdims))

    @property
    def group(self):
        if self._group:
            return self._group
        elements = [el for el in self if not el._auxiliary_component]
        values = {el.group for el in elements}
        types = {type(el) for el in elements}
        if values:
            group = list(values)[0]
            vtype = list(types)[0].__name__
        else:
            group, vtype = [], ''
        if len(values) == 1 and group != vtype:
            return group
        else:
            return type(self).__name__

    @group.setter
    def group(self, group):
        if not sanitize_identifier.allowable(group):
            raise ValueError("Supplied group %s contains invalid characters." %
                             group)
        else:
            self._group = group

    @property
    def label(self):
        if self._label:
            return self._label
        labels = {el.label for el in self
                  if not el._auxiliary_component}
        if len(labels) == 1:
            return list(labels)[0]
        else:
            return ''

    @label.setter
    def label(self, label):
        if not sanitize_identifier.allowable(label):
            raise ValueError("Supplied group %s contains invalid characters." %
                             label)
        self._label = label

    @property
    def ddims(self):
        dimensions = []
        dimension_names = []
        for el in self:
            for dim in el.dimensions():
                if dim.name not in dimension_names:
                    dimensions.append(dim)
                    dimension_names.append(dim.name)
        return dimensions

    @property
    def shape(self):
        raise NotImplementedError

class NdOverlay(UniformNdMapping, CompositeOverlay, Overlayable):
    """
    An NdOverlay allows a group of NdOverlay to be overlaid together. NdOverlay can
    be indexed out of an overlay and an overlay is an iterable that iterates
    over the contained layers.
    """

    kdims = param.List(default=[Dimension('Element')], constant=True, doc="""
        List of dimensions the NdOverlay can be indexed by.""")

    _deep_indexable = True

    def __init__(self, overlays=None, **params):
        super(NdOverlay, self).__init__(overlays, **params)

__all__ = list(set([_k for _k, _v in locals().items()
                    if isinstance(_v, type) and issubclass(_v, Dimensioned)])) + ['Overlayable']
<|endoftext|>


<commit_before>import argparse
import logging
import sys

import gym
from gym import wrappers
import os

def parse_args(args):
    parser = argparse.ArgumentParser(description=None)
    parser.add_argument('--env-id', dest='env_id', action='store',
                        default='CartPole-v0', help='Select the environment to run')
    parser.add_argument('--try-count', dest='episode_count', action='store',
                        default=10, type=int, help='Specify episode count')
    parser.add_argument('--output', dest='outdir', action='store',
                        default='/tmp/my-agent-results', help='Specify output directory')
    parser.add_argument('--api-key', dest='api_key', action='store',
                        default=None, help='Specify API key for OpenAI')
    parser.add_argument('--agent', dest='agent_name', action='store',
                        default='gymkit.agent.RandomAgent', help='Specify agent class name')
    parser.add_argument('--verbose', '-v', dest='verbose', action='store_true',
                        default=False, help='Display debug messages')
    return parser.parse_args(args=args)

def configure_logging(options):
    gym.undo_logger_setup()
    logger = logging.getLogger()
    formatter = logging.Formatter('[%(asctime)s] %(message)s')
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    if options.verbose:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    return logger

def create_gym_env(options):
    env = wrappers.Monitor(directory=options.outdir, force=True)(gym.make(options.env_id))
    env.seed(0)
    return env

def create_agent(env, options):
    sys.path.insert(0, os.getcwd())
    agent_values = options.agent_name.split('.')
    agent_module = '.'.join(agent_values[:-1])
    agent_name = agent_values[-1]
    mod = __import__(agent_module, globals(), locals(), fromlist=[agent_name])
    class_def = getattr(mod, agent_name)
    return class_def(env.action_space)

def main(args=None):
    options = parse_args(args)

    logger = configure_logging(options)
    env = create_gym_env(options)
    agent = create_agent(env, options)

    logger.info('Observation Space: %s', str(env.observation_space))
    logger.info('Action Space: %s', str(env.action_space))

    # environment loop
    for i in range(options.episode_count):
        reward = 0
        done = False
        ob = env.reset()
        with agent:
            logger.info('Episode %d', i + 1)
            while not done:
                env.render()
                logger.debug('O: %s', str(ob))
                action = agent.act(ob)
                logger.debug('A: %s', str(action))
                ob, reward, done, info = env.step(action)
                agent.fit(ob, reward, done, info)

    env.close()

    if options.api_key is not None:
        logger.info('Upload results from %s', options.outdir)
        gym.upload(options.outdir, api_key=options.api_key)

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>state checks for env and agent<commit_after>import argparse
import logging
import os
import sys

from gym import envs
from gym import wrappers
import gym

def parse_args(args):
    parser = argparse.ArgumentParser(description=None)
    parser.add_argument('--env-id', dest='env_id', action='store',
                        default='CartPole-v0', help='Select the environment to run')
    parser.add_argument('--try-count', dest='episode_count', action='store',
                        default=10, type=int, help='Specify episode count')
    parser.add_argument('--output', dest='outdir', action='store',
                        default='/tmp/my-agent-results', help='Specify output directory')
    parser.add_argument('--api-key', dest='api_key', action='store',
                        default=None, help='Specify API key for OpenAI')
    parser.add_argument('--agent', dest='agent_name', action='store',
                        default='gymkit.agent.RandomAgent', help='Specify agent class name')
    parser.add_argument('--verbose', '-v', dest='verbose', action='store_true',
                        default=False, help='Display debug messages')
    return parser.parse_args(args=args)

def configure_logging(options):
    gym.undo_logger_setup()
    logger = logging.getLogger()
    formatter = logging.Formatter('[%(asctime)s] %(message)s')
    handler = logging.StreamHandler(sys.stderr)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    if options.verbose:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)
    return logger

def create_gym_env(options):
    env = wrappers.Monitor(directory=options.outdir, force=True)(gym.make(options.env_id))
    env.seed(0)
    return env

def create_agent(env, options):
    sys.path.insert(0, os.getcwd())
    agent_values = options.agent_name.split('.')
    agent_module = '.'.join(agent_values[:-1])
    agent_name = agent_values[-1]
    mod = __import__(agent_module, globals(), locals(), fromlist=[agent_name])
    class_def = getattr(mod, agent_name)
    return class_def(env.action_space)

def main(args=None):
    options = parse_args(args)

    logger = configure_logging(options)

    try:
        env = create_gym_env(options)
    except:
        if logger.isEnabledFor(logging.DEBUG):
            logger.exception('Unknown environment %s.', options.env_id)
        else:
            env_ids = [x.id for x in envs.registry.all()]
            env_ids.sort()
            logger.error('Unknown environment %s. Available environments are:\n%s',
                         options.env_id,
                         '\n'.join(env_ids))
        return 1

    try:
        agent = create_agent(env, options)
    except:
        if logger.isEnabledFor(logging.DEBUG):
            logger.exception('Agent %s is not found.', options.agent_name)
        else:
            logger.error('Agent %s is not found.', options.agent_name)
        env.close()
        return 1

    logger.info('Observation Space: %s', str(env.observation_space))
    logger.info('Action Space: %s', str(env.action_space))

    # environment loop
    for i in range(options.episode_count):
        reward = 0
        done = False
        ob = env.reset()
        with agent:
            logger.info('Episode %d', i + 1)
            while not done:
                env.render()
                logger.debug('O: %s', str(ob))
                action = agent.act(ob)
                logger.debug('A: %s', str(action))
                ob, reward, done, info = env.step(action)
                agent.fit(ob, reward, done, info)

    env.close()

    if options.api_key is not None:
        logger.info('Upload results from %s', options.outdir)
        gym.upload(options.outdir, api_key=options.api_key)

    return 0

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


<commit_before>import itertools

import param
import numpy as np

from ..core import Dataset, OrderedDict
from ..core.operation import ElementOperation
from ..core.util import pd, is_nan, sort_topologically, cartesian_product

try:
    import dask
except:
    dask = None

try:
    import xarray as xr
except:
    xr = None

def toarray(v, index_value=False):
    """
    Interface helper function to turn dask Arrays into numpy arrays as
    necessary. If index_value is True, a value is returned instead of
    an array holding a single value.
    """
    if dask and isinstance(v, dask.array.Array):
        arr =  v.compute()
        return arr[()] if index_value else arr
    else:
        return v

def compute_edges(edges):
    """
    Computes edges from a number of bin centers,
    throwing an exception if the edges are not
    evenly spaced.
    """
    widths = np.diff(edges)
    if np.allclose(widths, widths[0]):
        width = widths[0]
    else:
        raise ValueError('Centered bins have to be of equal width.')
    edges -= width/2.
    return np.concatenate([edges, [edges[-1]+width]])

def reduce_fn(x):
    """
    Aggregation function to get the first non-zero value.
    """
    values = x.values if pd and isinstance(x, pd.Series) else x
    for v in values:
        if not is_nan(v):
            return v
    return np.NaN

class categorical_aggregate2d(ElementOperation):
    """
    Generates a gridded Dataset of 2D aggregate arrays indexed by the
    first two dimensions of the passed Element, turning all remaining
    dimensions into value dimensions. The key dimensions of the
    gridded array are treated as categorical indices. Useful for data
    indexed by two independent categorical variables such as a table
    of population values indexed by country and year. Data that is
    indexed by continuous dimensions should be binned before
    aggregation. The aggregation will retain the global sorting order
    of both dimensions.

    >> table = Table([('USA', 2000, 282.2), ('UK', 2005, 58.89)],
                     kdims=['Country', 'Year'], vdims=['Population'])
    >> categorical_aggregate2d(table)
    Dataset({'Country': ['USA', 'UK'], 'Year': [2000, 2005],
             'Population': [[ 282.2 , np.NaN], [np.NaN,   58.89]]},
            kdims=['Country', 'Year'], vdims=['Population'])
    """

    datatype = param.List(['xarray', 'grid'] if xr else ['grid'], doc="""
        The grid interface types to use when constructing the gridded Dataset.""")

    def _process(self, obj, key=None):
        """
        Generates a categorical 2D aggregate by inserting NaNs at all
        cross-product locations that do not already have a value assigned.
        Returns a 2D gridded Dataset object.
        """
        if isinstance(obj, Dataset) and obj.interface.gridded:
            return obj
        elif obj.ndims > 2:
            raise ValueError("Cannot aggregate more than two dimensions")
        elif len(obj.dimensions()) < 3:
            raise ValueError("Must have at two dimensions to aggregate over"
                             "and one value dimension to aggregate on.")

        dim_labels = obj.dimensions(label=True)
        dims = obj.dimensions()
        kdims, vdims = dims[:2], dims[2:]
        xdim, ydim = dim_labels[:2]
        nvdims = len(dims) - 2
        d1keys = obj.dimension_values(xdim, False)
        d2keys = obj.dimension_values(ydim, False)
        shape = (len(d2keys), len(d1keys))
        nsamples = np.product(shape)

        # Determine global orderings of y-values using topological sort
        grouped = obj.groupby(xdim, container_type=OrderedDict,
                              group_type=Dataset).values()
        orderings = OrderedDict()
        is_sorted = np.array_equal(np.sort(d1keys), d1keys)
        for group in grouped:
            vals = group.dimension_values(ydim)
            if len(vals) == 1:
                orderings[vals[0]] = []
            else:
                is_sorted &= np.array_equal(np.sort(vals), vals)
                for i in range(len(vals)-1):
                    p1, p2 = vals[i:i+2]
                    orderings[p1] = [p2]
        if is_sorted:
            d2keys = np.sort(d2keys)
        else:
            d2keys = list(itertools.chain(*sort_topologically(orderings)))

        # Pad data with NaNs
        ys, xs = cartesian_product([d2keys, d1keys])
        data = {xdim: xs.flatten(), ydim: ys.flatten()}
        for vdim in vdims:
            values = np.empty(nsamples)
            values[:] = np.NaN
            data[vdim.name] = values
        dtype = 'dataframe' if pd else 'dictionary'
        dense_data = Dataset(data, kdims=obj.kdims, vdims=obj.vdims, datatype=[dtype])
        concat_data = obj.interface.concatenate([dense_data, Dataset(obj)], datatype=dtype)
        agg = concat_data.reindex([xdim, ydim]).aggregate([xdim, ydim], reduce_fn)

        # Convert data to a gridded dataset
        grid_data = {xdim: d1keys, ydim: d2keys}
        for vdim in vdims:
            grid_data[vdim.name] = agg.dimension_values(vdim).reshape(shape)
        grid_type = 'xarray' if xr else 'grid'
        return agg.clone(grid_data, datatype=self.p.datatype)

<commit_msg>Removed unused variable<commit_after>import itertools

import param
import numpy as np

from ..core import Dataset, OrderedDict
from ..core.operation import ElementOperation
from ..core.util import pd, is_nan, sort_topologically, cartesian_product

try:
    import dask
except:
    dask = None

try:
    import xarray as xr
except:
    xr = None

def toarray(v, index_value=False):
    """
    Interface helper function to turn dask Arrays into numpy arrays as
    necessary. If index_value is True, a value is returned instead of
    an array holding a single value.
    """
    if dask and isinstance(v, dask.array.Array):
        arr =  v.compute()
        return arr[()] if index_value else arr
    else:
        return v

def compute_edges(edges):
    """
    Computes edges from a number of bin centers,
    throwing an exception if the edges are not
    evenly spaced.
    """
    widths = np.diff(edges)
    if np.allclose(widths, widths[0]):
        width = widths[0]
    else:
        raise ValueError('Centered bins have to be of equal width.')
    edges -= width/2.
    return np.concatenate([edges, [edges[-1]+width]])

def reduce_fn(x):
    """
    Aggregation function to get the first non-zero value.
    """
    values = x.values if pd and isinstance(x, pd.Series) else x
    for v in values:
        if not is_nan(v):
            return v
    return np.NaN

class categorical_aggregate2d(ElementOperation):
    """
    Generates a gridded Dataset of 2D aggregate arrays indexed by the
    first two dimensions of the passed Element, turning all remaining
    dimensions into value dimensions. The key dimensions of the
    gridded array are treated as categorical indices. Useful for data
    indexed by two independent categorical variables such as a table
    of population values indexed by country and year. Data that is
    indexed by continuous dimensions should be binned before
    aggregation. The aggregation will retain the global sorting order
    of both dimensions.

    >> table = Table([('USA', 2000, 282.2), ('UK', 2005, 58.89)],
                     kdims=['Country', 'Year'], vdims=['Population'])
    >> categorical_aggregate2d(table)
    Dataset({'Country': ['USA', 'UK'], 'Year': [2000, 2005],
             'Population': [[ 282.2 , np.NaN], [np.NaN,   58.89]]},
            kdims=['Country', 'Year'], vdims=['Population'])
    """

    datatype = param.List(['xarray', 'grid'] if xr else ['grid'], doc="""
        The grid interface types to use when constructing the gridded Dataset.""")

    def _process(self, obj, key=None):
        """
        Generates a categorical 2D aggregate by inserting NaNs at all
        cross-product locations that do not already have a value assigned.
        Returns a 2D gridded Dataset object.
        """
        if isinstance(obj, Dataset) and obj.interface.gridded:
            return obj
        elif obj.ndims > 2:
            raise ValueError("Cannot aggregate more than two dimensions")
        elif len(obj.dimensions()) < 3:
            raise ValueError("Must have at two dimensions to aggregate over"
                             "and one value dimension to aggregate on.")

        dim_labels = obj.dimensions(label=True)
        dims = obj.dimensions()
        kdims, vdims = dims[:2], dims[2:]
        xdim, ydim = dim_labels[:2]
        nvdims = len(dims) - 2
        d1keys = obj.dimension_values(xdim, False)
        d2keys = obj.dimension_values(ydim, False)
        shape = (len(d2keys), len(d1keys))
        nsamples = np.product(shape)

        # Determine global orderings of y-values using topological sort
        grouped = obj.groupby(xdim, container_type=OrderedDict,
                              group_type=Dataset).values()
        orderings = OrderedDict()
        is_sorted = np.array_equal(np.sort(d1keys), d1keys)
        for group in grouped:
            vals = group.dimension_values(ydim)
            if len(vals) == 1:
                orderings[vals[0]] = []
            else:
                is_sorted &= np.array_equal(np.sort(vals), vals)
                for i in range(len(vals)-1):
                    p1, p2 = vals[i:i+2]
                    orderings[p1] = [p2]
        if is_sorted:
            d2keys = np.sort(d2keys)
        else:
            d2keys = list(itertools.chain(*sort_topologically(orderings)))

        # Pad data with NaNs
        ys, xs = cartesian_product([d2keys, d1keys])
        data = {xdim: xs.flatten(), ydim: ys.flatten()}
        for vdim in vdims:
            values = np.empty(nsamples)
            values[:] = np.NaN
            data[vdim.name] = values
        dtype = 'dataframe' if pd else 'dictionary'
        dense_data = Dataset(data, kdims=obj.kdims, vdims=obj.vdims, datatype=[dtype])
        concat_data = obj.interface.concatenate([dense_data, Dataset(obj)], datatype=dtype)
        agg = concat_data.reindex([xdim, ydim]).aggregate([xdim, ydim], reduce_fn)

        # Convert data to a gridded dataset
        grid_data = {xdim: d1keys, ydim: d2keys}
        for vdim in vdims:
            grid_data[vdim.name] = agg.dimension_values(vdim).reshape(shape)
        return agg.clone(grid_data, datatype=self.p.datatype)

<|endoftext|>


<commit_before># vim: tabstop=4 shiftwidth=4 softtabstop=4

#
#    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.

"""
/stack endpoint for heat v1 API
"""
import dbus
import httplib
import json
import libxml2
import logging
import sys
import urlparse

import webob
from webob.exc import (HTTPNotFound,
                       HTTPConflict,
                       HTTPBadRequest)

from heat.common import exception
from heat.common import utils
from heat.common import wsgi

logger = logging.getLogger('heat.api.v1.stacks')

stack_db = {}

class Json2CapeXml:
    def __init__(self, template, stack_name):

        self.t = template
        self.parms = self.t['Parameters']
        self.maps = self.t['Mappings']
        self.res = {}

        self.parms['AWS::Region'] = {"Description" : "AWS Regions", "Type" : "String", "Default" : "ap-southeast-1",
              "AllowedValues" : ["us-east-1","us-west-1","us-west-2","sa-east-1","eu-west-1","ap-southeast-1","ap-northeast-1"],
              "ConstraintDescription" : "must be a valid EC2 instance type." }

        # expected user parameters
        self.parms['AWS::StackName'] = {'Default': stack_name}
        self.parms['KeyName'] = {'Default': 'harry-45-5-34-5'}

        for r in self.t['Resources']:
            # fake resource instance references
            self.parms[r] = {'Default': utils.generate_uuid()}

        self.resolve_static_refs(self.t['Resources'])
        self.resolve_find_in_map(self.t['Resources'])
        #self.resolve_attributes(self.t['Resources'])
        self.resolve_joins(self.t['Resources'])
        self.resolve_base64(self.t['Resources'])
        #print json.dumps(self.t['Resources'], indent=2)

    def convert_and_write(self):

        name = self.parms['AWS::StackName']['Default']

        doc = libxml2.newDoc("1.0")
        dep = doc.newChild(None, "deployable", None)
        dep.setProp("name", name)
        dep.setProp("uuid", 'bogus')
        dep.setProp("monitor", 'active')
        dep.setProp("username", 'nobody-yet')
        n_asses = dep.newChild(None, "assemblies", None)

        for r in self.t['Resources']:
            type = self.t['Resources'][r]['Type']
            if type != 'AWS::EC2::Instance':
                print 'ignoring Resource %s (%s)' % (r, type)
                continue

            n_ass = n_asses.newChild(None, 'assembly', None)
            n_ass.setProp("name", r)
            n_ass.setProp("uuid", self.parms[r]['Default'])
            props = self.t['Resources'][r]['Properties']
            for p in props:
                if p == 'ImageId':
                    n_ass.setProp("image_name", props[p])
                elif p == 'UserData':
                    new_script = []
                    script_lines = props[p].split('\n')
                    for l in script_lines:
                        if '#!/' in l:
                            new_script.append(l)
                            self.insert_package_and_services(self.t['Resources'][r], new_script)
                        else:
                            new_script.append(l)

                    startup = n_ass.newChild(None, 'startup', '\n'.join(new_script))

            con = self.t['Resources'][r]['Metadata']["AWS::CloudFormation::Init"]['config']

            n_services = n_ass.newChild(None, 'services', None)
            for st in con['services']:
                for s in con['services'][st]:
                    n_service = n_services.newChild(None, 'service', None)
                    n_service.setProp("name", '%s_%s' % (r, s))
                    n_service.setProp("type", s)
                    n_service.setProp("provider", 'pacemaker')
                    n_service.setProp("class", 'lsb')
                    n_service.setProp("monitor_interval", '30s')
                    n_service.setProp("escalation_period", '1000')
                    n_service.setProp("escalation_failures", '3')

        filename = '/var/run/%s.xml' % name
        open(filename, 'w').write(doc.serialize(None, 1))
        doc.freeDoc()

    def insert_package_and_services(self, r, new_script):

        con = r['Metadata']["AWS::CloudFormation::Init"]['config']

        for pt in con['packages']:
            if pt == 'yum':
                for p in con['packages']['yum']:
                    new_script.append('yum install -y %s' % p)
        for st in con['services']:
            if st == 'systemd':
                for s in con['services']['systemd']:
                    v = con['services']['systemd'][s]
                    if v['enabled'] == 'true':
                        new_script.append('systemctl enable %s.service' % s)
                    if v['ensureRunning'] == 'true':
                        new_script.append('systemctl start %s.service' % s)
            elif st == 'sysvinit':
                for s in con['services']['sysvinit']:
                    v = con['services']['systemd'][s]
                    if v['enabled'] == 'true':
                        new_script.append('chkconfig %s on' % s)
                    if v['ensureRunning'] == 'true':
                        new_script.append('/etc/init.d/start %s' % s)

    def resolve_static_refs(self, s):
        '''
            looking for { "Ref": "str" }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Ref' and isinstance(s[i], (basestring, unicode)) and \
                      self.parms.has_key(s[i]):
                    if self.parms[s[i]] == None:
                        print 'None Ref: %s' % str(s[i])
                    elif self.parms[s[i]].has_key('Default'):
                        # note the "ref: values" are in a dict of
                        # size one, so return is fine.
                        #print 'Ref: %s == %s' % (s[i], self.parms[s[i]]['Default'])
                        return self.parms[s[i]]['Default']
                    else:
                        print 'missing Ref: %s' % str(s[i])
                else:
                    s[i] = self.resolve_static_refs(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                #print 'resolve_static_refs %d %s' % (index, item)
                s[index] = self.resolve_static_refs(item)
        return s

    def resolve_find_in_map(self, s):
        '''
            looking for { "Ref": "str" }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::FindInMap':
                    obj = self.maps
                    if isinstance(s[i], list):
                        #print 'map list: %s' % s[i]
                        for index, item in enumerate(s[i]):
                            if isinstance(item, dict):
                                item = self.resolve_find_in_map(item)
                                #print 'map item dict: %s' % (item)
                            else:
                                pass
                                #print 'map item str: %s' % (item)
                            obj = obj[item]
                    else:
                        obj = obj[s[i]]
                    return obj
                else:
                    s[i] = self.resolve_find_in_map(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_find_in_map(item)
        return s

    def resolve_joins(self, s):
        '''
            looking for { "Fn::join": [] }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::Join':
                    return s[i][0].join(s[i][1])
                else:
                    s[i] = self.resolve_joins(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_joins(item)
        return s

    def resolve_base64(self, s):
        '''
            looking for { "Fn::join": [] }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::Base64':
                    return s[i]
                else:
                    s[i] = self.resolve_base64(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_base64(item)
        return s

def systemctl(method, name, instance=None):

    bus = dbus.SystemBus()

    sysd = bus.get_object('org.freedesktop.systemd1',
                         '/org/freedesktop/systemd1')

    actual_method = ''
    if method == 'start':
        actual_method = 'StartUnit'
    elif method == 'stop':
        actual_method = 'StopUnit'
    else:
        raise

    m = sysd.get_dbus_method(actual_method, 'org.freedesktop.systemd1.Manager')

    if instance == None:
        service = '%s.service' % (name)
    else:
        service = '%s@%s.service' % (name, instance)

    result = m(service, 'replace')
    return result

class StackController(object):

    """
    WSGI controller for stacks resource in heat v1 API

    """

    def __init__(self, options):
        self.options = options
        self.stack_id = 1

    def list(self, req):
        """
        Returns the following information for all stacks:
        """
        res = {'ListStacksResponse': {'ListStacksResult': {'StackSummaries': [] } } }
        summaries = res['ListStacksResponse']['ListStacksResult']['StackSummaries']
        for s in stack_db:
            mem = {}
            mem['StackId'] = stack_db[s]['StackId']
            mem['StackName'] = s
            mem['CreationTime'] = 'now'
            try:
                mem['TemplateDescription'] = stack_db[s]['Description']
                mem['StackStatus'] = stack_db[s]['StackStatus']
            except:
                mem['TemplateDescription'] = 'No description'
                mem['StackStatus'] = 'unknown'
            summaries.append(mem)

        return res

    def describe(self, req):

        stack_name = None
        if req.params.has_key('StackName'):
            stack_name = req.params['StackName']
            if not stack_db.has_key(stack_name):
                msg = _("Stack does not exist with that name.")
                return webob.exc.HTTPNotFound(msg)

        res = {'DescribeStacksResult': {'Stacks': [] } }
        summaries = res['DescribeStacksResult']['Stacks']
        for s in stack_db:
            if stack_name is None or s == stack_name:
                mem = {}
                mem['StackId'] = stack_db[s]['StackId']
                mem['StackStatus'] = stack_db[s]['StackStatus']
                mem['StackName'] = s
                mem['CreationTime'] = 'now'
                mem['DisableRollback'] = 'false'
                mem['Outputs'] = '[]'
                summaries.append(mem)

        return res

    def _get_template(self, req):
        if req.params.has_key('TemplateBody'):
            logger.info('TemplateBody ...')
            return req.params['TemplateBody']
        elif req.params.has_key('TemplateUrl'):
            logger.info('TemplateUrl %s' % req.params['TemplateUrl'])
            url = urlparse.urlparse(req.params['TemplateUrl'])
            if url.scheme == 'https':
                conn = httplib.HTTPSConnection(url.netloc)
            else:
                conn = httplib.HTTPConnection(url.netloc)
            conn.request("GET", url.path)
            r1 = conn.getresponse()
            logger.info('status %d' % r1.status)
            if r1.status == 200:
                data = r1.read()
                conn.close()
            else:
                data = None
            return data

        return None

    def _apply_user_parameters(self, req, stack):
        # TODO
        pass

    def create(self, req):
        """
        :param req: The WSGI/Webob Request object

        :raises HttpBadRequest if not template is given
        :raises HttpConflict if object already exists
        """
        if stack_db.has_key(req.params['StackName']):
            msg = _("Stack already exists with that name.")
            return webob.exc.HTTPConflict(msg)

        templ = self._get_template(req)
        if templ is None:
            msg = _("TemplateBody or TemplateUrl were not given.")
            return webob.exc.HTTPBadRequest(explanation=msg)

        stack = json.loads(templ)
        my_id = '%s-%d' % (req.params['StackName'], self.stack_id)
        self.stack_id = self.stack_id + 1
        stack['StackId'] = my_id
        stack['StackStatus'] = 'CREATE_COMPLETE'
        self._apply_user_parameters(req, stack)
        stack_db[req.params['StackName']] = stack

        cape_transformer = Json2CapeXml(stack, req.params['StackName'])
        cape_transformer.convert_and_write()

        systemctl('start', 'pcloud-cape-sshd', req.params['StackName'])

        return {'CreateStackResult': {'StackId': my_id}}

    def update(self, req):
        """
        :param req: The WSGI/Webob Request object

        :raises HttpNotFound if object is not available
        """
        if not stack_db.has_key(req.params['StackName']):
            msg = _("Stack does not exist with that name.")
            return webob.exc.HTTPNotFound(msg)

        stack = stack_db[req.params['StackName']]
        my_id = stack['StackId']
        templ = self._get_template(req)
        if templ:
            stack = json.loads(templ)
            stack['StackId'] = my_id
            stack_db[req.params['StackName']] = stack

        self._apply_user_parameters(req, stack)
        stack['StackStatus'] = 'UPDATE_COMPLETE'

        return {'UpdateStackResult': {'StackId': my_id}}

    def delete(self, req):
        """
        Deletes the object and all its resources

        :param req: The WSGI/Webob Request object

        :raises HttpBadRequest if the request is invalid
        :raises HttpNotFound if object is not available
        :raises HttpNotAuthorized if object is not
                deleteable by the requesting user
        """
        logger.info('in delete %s ' % req.params['StackName'])
        if not stack_db.has_key(req.params['StackName']):
            msg = _("Stack does not exist with that name.")
            return webob.exc.HTTPNotFound(msg)

        del stack_db[req.params['StackName']]

        systemctl('stop', 'pcloud-cape-sshd', req.params['StackName'])

def create_resource(options):
    """Stacks resource factory method"""
    deserializer = wsgi.JSONRequestDeserializer()
    serializer = wsgi.JSONResponseSerializer()
    return wsgi.Resource(StackController(options), deserializer, serializer)
<commit_msg>Don't crash if there is no "AWS::CloudFormation::Init" section.<commit_after># vim: tabstop=4 shiftwidth=4 softtabstop=4

#
#    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.

"""
/stack endpoint for heat v1 API
"""
import dbus
import httplib
import json
import libxml2
import logging
import sys
import urlparse

import webob
from webob.exc import (HTTPNotFound,
                       HTTPConflict,
                       HTTPBadRequest)

from heat.common import exception
from heat.common import utils
from heat.common import wsgi

logger = logging.getLogger('heat.api.v1.stacks')

stack_db = {}

class Json2CapeXml:
    def __init__(self, template, stack_name):

        self.t = template
        self.parms = self.t['Parameters']
        self.maps = self.t['Mappings']
        self.res = {}

        self.parms['AWS::Region'] = {"Description" : "AWS Regions", "Type" : "String", "Default" : "ap-southeast-1",
              "AllowedValues" : ["us-east-1","us-west-1","us-west-2","sa-east-1","eu-west-1","ap-southeast-1","ap-northeast-1"],
              "ConstraintDescription" : "must be a valid EC2 instance type." }

        # expected user parameters
        self.parms['AWS::StackName'] = {'Default': stack_name}
        self.parms['KeyName'] = {'Default': 'harry-45-5-34-5'}

        for r in self.t['Resources']:
            # fake resource instance references
            self.parms[r] = {'Default': utils.generate_uuid()}

        self.resolve_static_refs(self.t['Resources'])
        self.resolve_find_in_map(self.t['Resources'])
        #self.resolve_attributes(self.t['Resources'])
        self.resolve_joins(self.t['Resources'])
        self.resolve_base64(self.t['Resources'])
        #print json.dumps(self.t['Resources'], indent=2)

    def convert_and_write(self):

        name = self.parms['AWS::StackName']['Default']

        doc = libxml2.newDoc("1.0")
        dep = doc.newChild(None, "deployable", None)
        dep.setProp("name", name)
        dep.setProp("uuid", 'bogus')
        dep.setProp("monitor", 'active')
        dep.setProp("username", 'nobody-yet')
        n_asses = dep.newChild(None, "assemblies", None)

        for r in self.t['Resources']:
            type = self.t['Resources'][r]['Type']
            if type != 'AWS::EC2::Instance':
                print 'ignoring Resource %s (%s)' % (r, type)
                continue

            n_ass = n_asses.newChild(None, 'assembly', None)
            n_ass.setProp("name", r)
            n_ass.setProp("uuid", self.parms[r]['Default'])
            props = self.t['Resources'][r]['Properties']
            for p in props:
                if p == 'ImageId':
                    n_ass.setProp("image_name", props[p])
                elif p == 'UserData':
                    new_script = []
                    script_lines = props[p].split('\n')
                    for l in script_lines:
                        if '#!/' in l:
                            new_script.append(l)
                            self.insert_package_and_services(self.t['Resources'][r], new_script)
                        else:
                            new_script.append(l)

                    startup = n_ass.newChild(None, 'startup', '\n'.join(new_script))

            try:
                con = self.t['Resources'][r]['Metadata']["AWS::CloudFormation::Init"]['config']
                n_services = n_ass.newChild(None, 'services', None)
                for st in con['services']:
                    for s in con['services'][st]:
                        n_service = n_services.newChild(None, 'service', None)
                        n_service.setProp("name", '%s_%s' % (r, s))
                        n_service.setProp("type", s)
                        n_service.setProp("provider", 'pacemaker')
                        n_service.setProp("class", 'lsb')
                        n_service.setProp("monitor_interval", '30s')
                        n_service.setProp("escalation_period", '1000')
                        n_service.setProp("escalation_failures", '3')
            except KeyError as e:
                # if there is no config then no services.
                pass

        filename = '/var/run/%s.xml' % name
        open(filename, 'w').write(doc.serialize(None, 1))
        doc.freeDoc()

    def insert_package_and_services(self, r, new_script):

        try:
            con = r['Metadata']["AWS::CloudFormation::Init"]['config']
        except KeyError as e:
            return

        for pt in con['packages']:
            if pt == 'yum':
                for p in con['packages']['yum']:
                    new_script.append('yum install -y %s' % p)
        for st in con['services']:
            if st == 'systemd':
                for s in con['services']['systemd']:
                    v = con['services']['systemd'][s]
                    if v['enabled'] == 'true':
                        new_script.append('systemctl enable %s.service' % s)
                    if v['ensureRunning'] == 'true':
                        new_script.append('systemctl start %s.service' % s)
            elif st == 'sysvinit':
                for s in con['services']['sysvinit']:
                    v = con['services']['systemd'][s]
                    if v['enabled'] == 'true':
                        new_script.append('chkconfig %s on' % s)
                    if v['ensureRunning'] == 'true':
                        new_script.append('/etc/init.d/start %s' % s)

    def resolve_static_refs(self, s):
        '''
            looking for { "Ref": "str" }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Ref' and isinstance(s[i], (basestring, unicode)) and \
                      self.parms.has_key(s[i]):
                    if self.parms[s[i]] == None:
                        print 'None Ref: %s' % str(s[i])
                    elif self.parms[s[i]].has_key('Default'):
                        # note the "ref: values" are in a dict of
                        # size one, so return is fine.
                        #print 'Ref: %s == %s' % (s[i], self.parms[s[i]]['Default'])
                        return self.parms[s[i]]['Default']
                    else:
                        print 'missing Ref: %s' % str(s[i])
                else:
                    s[i] = self.resolve_static_refs(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                #print 'resolve_static_refs %d %s' % (index, item)
                s[index] = self.resolve_static_refs(item)
        return s

    def resolve_find_in_map(self, s):
        '''
            looking for { "Ref": "str" }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::FindInMap':
                    obj = self.maps
                    if isinstance(s[i], list):
                        #print 'map list: %s' % s[i]
                        for index, item in enumerate(s[i]):
                            if isinstance(item, dict):
                                item = self.resolve_find_in_map(item)
                                #print 'map item dict: %s' % (item)
                            else:
                                pass
                                #print 'map item str: %s' % (item)
                            obj = obj[item]
                    else:
                        obj = obj[s[i]]
                    return obj
                else:
                    s[i] = self.resolve_find_in_map(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_find_in_map(item)
        return s

    def resolve_joins(self, s):
        '''
            looking for { "Fn::join": [] }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::Join':
                    return s[i][0].join(s[i][1])
                else:
                    s[i] = self.resolve_joins(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_joins(item)
        return s

    def resolve_base64(self, s):
        '''
            looking for { "Fn::join": [] }
        '''
        if isinstance(s, dict):
            for i in s:
                if i == 'Fn::Base64':
                    return s[i]
                else:
                    s[i] = self.resolve_base64(s[i])
        elif isinstance(s, list):
            for index, item in enumerate(s):
                s[index] = self.resolve_base64(item)
        return s

def systemctl(method, name, instance=None):

    bus = dbus.SystemBus()

    sysd = bus.get_object('org.freedesktop.systemd1',
                         '/org/freedesktop/systemd1')

    actual_method = ''
    if method == 'start':
        actual_method = 'StartUnit'
    elif method == 'stop':
        actual_method = 'StopUnit'
    else:
        raise

    m = sysd.get_dbus_method(actual_method, 'org.freedesktop.systemd1.Manager')

    if instance == None:
        service = '%s.service' % (name)
    else:
        service = '%s@%s.service' % (name, instance)

    result = m(service, 'replace')
    return result

class StackController(object):

    """
    WSGI controller for stacks resource in heat v1 API

    """

    def __init__(self, options):
        self.options = options
        self.stack_id = 1

    def list(self, req):
        """
        Returns the following information for all stacks:
        """
        res = {'ListStacksResponse': {'ListStacksResult': {'StackSummaries': [] } } }
        summaries = res['ListStacksResponse']['ListStacksResult']['StackSummaries']
        for s in stack_db:
            mem = {}
            mem['StackId'] = stack_db[s]['StackId']
            mem['StackName'] = s
            mem['CreationTime'] = 'now'
            try:
                mem['TemplateDescription'] = stack_db[s]['Description']
                mem['StackStatus'] = stack_db[s]['StackStatus']
            except:
                mem['TemplateDescription'] = 'No description'
                mem['StackStatus'] = 'unknown'
            summaries.append(mem)

        return res

    def describe(self, req):

        stack_name = None
        if req.params.has_key('StackName'):
            stack_name = req.params['StackName']
            if not stack_db.has_key(stack_name):
                msg = _("Stack does not exist with that name.")
                return webob.exc.HTTPNotFound(msg)

        res = {'DescribeStacksResult': {'Stacks': [] } }
        summaries = res['DescribeStacksResult']['Stacks']
        for s in stack_db:
            if stack_name is None or s == stack_name:
                mem = {}
                mem['StackId'] = stack_db[s]['StackId']
                mem['StackStatus'] = stack_db[s]['StackStatus']
                mem['StackName'] = s
                mem['CreationTime'] = 'now'
                mem['DisableRollback'] = 'false'
                mem['Outputs'] = '[]'
                summaries.append(mem)

        return res

    def _get_template(self, req):
        if req.params.has_key('TemplateBody'):
            logger.info('TemplateBody ...')
            return req.params['TemplateBody']
        elif req.params.has_key('TemplateUrl'):
            logger.info('TemplateUrl %s' % req.params['TemplateUrl'])
            url = urlparse.urlparse(req.params['TemplateUrl'])
            if url.scheme == 'https':
                conn = httplib.HTTPSConnection(url.netloc)
            else:
                conn = httplib.HTTPConnection(url.netloc)
            conn.request("GET", url.path)
            r1 = conn.getresponse()
            logger.info('status %d' % r1.status)
            if r1.status == 200:
                data = r1.read()
                conn.close()
            else:
                data = None
            return data

        return None

    def _apply_user_parameters(self, req, stack):
        # TODO
        pass

    def create(self, req):
        """
        :param req: The WSGI/Webob Request object

        :raises HttpBadRequest if not template is given
        :raises HttpConflict if object already exists
        """
        if stack_db.has_key(req.params['StackName']):
            msg = _("Stack already exists with that name.")
            return webob.exc.HTTPConflict(msg)

        templ = self._get_template(req)
        if templ is None:
            msg = _("TemplateBody or TemplateUrl were not given.")
            return webob.exc.HTTPBadRequest(explanation=msg)

        stack = json.loads(templ)
        my_id = '%s-%d' % (req.params['StackName'], self.stack_id)
        self.stack_id = self.stack_id + 1
        stack['StackId'] = my_id
        stack['StackStatus'] = 'CREATE_COMPLETE'
        self._apply_user_parameters(req, stack)
        stack_db[req.params['StackName']] = stack

        cape_transformer = Json2CapeXml(stack, req.params['StackName'])
        cape_transformer.convert_and_write()

        systemctl('start', 'pcloud-cape-sshd', req.params['StackName'])

        return {'CreateStackResult': {'StackId': my_id}}

    def update(self, req):
        """
        :param req: The WSGI/Webob Request object

        :raises HttpNotFound if object is not available
        """
        if not stack_db.has_key(req.params['StackName']):
            msg = _("Stack does not exist with that name.")
            return webob.exc.HTTPNotFound(msg)

        stack = stack_db[req.params['StackName']]
        my_id = stack['StackId']
        templ = self._get_template(req)
        if templ:
            stack = json.loads(templ)
            stack['StackId'] = my_id
            stack_db[req.params['StackName']] = stack

        self._apply_user_parameters(req, stack)
        stack['StackStatus'] = 'UPDATE_COMPLETE'

        return {'UpdateStackResult': {'StackId': my_id}}

    def delete(self, req):
        """
        Deletes the object and all its resources

        :param req: The WSGI/Webob Request object

        :raises HttpBadRequest if the request is invalid
        :raises HttpNotFound if object is not available
        :raises HttpNotAuthorized if object is not
                deleteable by the requesting user
        """
        logger.info('in delete %s ' % req.params['StackName'])
        if not stack_db.has_key(req.params['StackName']):
            msg = _("Stack does not exist with that name.")
            return webob.exc.HTTPNotFound(msg)

        del stack_db[req.params['StackName']]

        systemctl('stop', 'pcloud-cape-sshd', req.params['StackName'])

def create_resource(options):
    """Stacks resource factory method"""
    deserializer = wsgi.JSONRequestDeserializer()
    serializer = wsgi.JSONResponseSerializer()
    return wsgi.Resource(StackController(options), deserializer, serializer)
<|endoftext|>