<commit_before>import numpy as np
import sys
import os
import subprocess as sp
import webbrowser as wb

# Global variables
iconst = 1 # Multiplier for individual array elements in integer and real arrays read by MODFLOW's U2DREL, U1DREL and U2DINT.
iprn = -1 # Printout flag. If >= 0 then array values read are printed in listing file.

class BaseModel(object):
    """
    MODFLOW based models base class
    """

    def __init__(self, modelname = 'modflowtest', namefile_ext = 'nam',
                 exe_name = 'mf2k.exe', model_ws = None):
        self.__name = modelname
        self.namefile_ext = namefile_ext
        self.namefile = self.__name + '.' + self.namefile_ext
        self.packagelist = []
        self.heading = ''
        self.exe_name = exe_name
        self.external_extension = 'ref'
        if model_ws is None: model_ws = os.getcwd()
        if not os.path.exists(model_ws):
            try:
                os.makedirs(model_ws)
            except:
                #print '\n%s not valid, workspace-folder was changed to %s\n' % (model_ws, os.getcwd())
                print '\n{0:s} not valid, workspace-folder was changed to {1:s}\n'.format(model_ws, os.getcwd())
                model_ws = os.getcwd()
        self.model_ws= model_ws
        self.cl_params = ''

    def set_exename(self, exe_name):
        self.exe_name = exe_name
        return

    def add_package(self, p):
        """
        Add a flopy package object to this model.
        """
        for pp in (self.packagelist):
            if pp.allowDuplicates:
                continue
            elif (isinstance(p, type(pp))):
                print '****Warning -- two packages of the same type: ',type(p),type(pp)
                print 'replacing existing Package...'
                pp = p
                return
        self.packagelist.append( p )

    def remove_package(self, pname):
        """
        Remove a package from this model.
        """
        for i,pp in enumerate(self.packagelist):
            if pname in pp.name:
                print 'removing Package: ',pp.name
                self.packagelist.pop(i)
                return
        raise StopIteration , 'Package name '+pname+' not found in Package list'

    def build_array_name(self,num,prefix):
       return self.external_path+prefix+'_'+str(num)+'.'+self.external_extension

    def assign_external(self,num,prefix):
        fname = self.build_array_name(num,prefix)
        unit = (self.next_ext_unit())
        self.external_fnames.append(fname)
        self.external_units.append(unit)
        self.external_binflag.append(False)
        return fname,unit

    def add_external(self, fname, unit, binflag=False):
        """
        Supports SWR usage and non-loaded existing external arrays
        """
        self.external_fnames.append(fname)
        self.external_units.append(unit)
        self.external_binflag.append(binflag)
        return

    def remove_external(self,fname=None,unit=None):
        if fname is not None:
            for i,e in enumerate(self.external_fnames):
                if fname in e:
                    self.external_fnames.pop(i)
                    self.external_units.pop(i)
                    self.external_binflag.pop(i)
        elif unit is not None:
            for i,u in enumerate(self.external_units):
                if u == unit:
                    self.external_fnames.pop(i)
                    self.external_units.pop(i)
                    self.external_binflag.pop(i)
        else:
            raise Exception,' either fname or unit must be passed to remove_external()'
        return

    def get_name_file_entries(self):
        s = ''
        for p in self.packagelist:
            for i in range(len(p.name)):
                #s = s + ('%s %3i %s %s\n' % (p.name[i], p.unit_number[i],
                #                             p.file_name[i],p.extra[i]))
                s = s + ('{0:s} {1:3d} {2:s} {3:s}\n'.format(p.name[i], p.unit_number[i],
                                             p.file_name[i],p.extra[i]))
        return s

    def get_package(self, name):
        for pp in (self.packagelist):
            if (pp.name[0].upper() == name.upper()):
                return pp
        return None

    def get_package_list(self):
        val = []
        for pp in (self.packagelist):
            val.append(pp.name[0].upper())
        return val

    def change_model_ws(self,new_pth=None):
        if new_pth is None:
            new_pth = os.getcwd()
        if not os.path.exists(new_pth):
            try:
                os.makedirs(new_pth)
            except:
                #print '\n%s not valid, workspace-folder was changed to %s\n' % (new_pth, os.getcwd())
                print '\n{0:s} not valid, workspace-folder was changed to {1:s}\n'.format(new_pth, os.getcwd())
                new_pth = os.getcwd()
        #--reset the model workspace
        self.model_ws = new_pth
        #--reset the paths for each package
        for pp in (self.packagelist):
            pp.fn_path = os.path.join(self.model_ws,pp.file_name[0])
        return None

    def run_model(self, pause=True, report=None):
        """
        Run the model.  This method will create and run a Windows batch file.

        Parameters
        ----------
        pause : boolean, optional
            Add a pause to the end of the batch file.
            (the default is False)
        report : string, optional
            Name of file to store stdout. (default is None).
        """
        batch_file_name = os.path.join(self.model_ws, 'run.bat')
        #error_message = ('Model executable %s not found!' % self.exe_name)
        error_message = 'Model executable {0:s} not found!'.format(self.exe_name)
        assert os.path.exists(self.exe_name), error_message

        error_message = ('Name file %s not found!' % self.namefile)
        fn_path = os.path.join(self.model_ws, self.namefile)
        assert os.path.exists(fn_path), error_message

        # Create a batch file to call code so that window remains open in case of error messages
        f = open(batch_file_name, 'w')
        #f.write('@ECHO Calling %s with %s\n' % (self.exe_name, self.namefile))
        f.write('@ECHO Calling {0:s} with {1:s}\n'.format(self.exe_name, self.namefile))
        #f.write('%s %s %s\n' % (self.exe_name, self.namefile, self.cl_params))
        f.write('{0:s} {1:s} {2:s}\n'.format(self.exe_name, self.namefile, self.cl_params))
        if (pause):
           f.write('@PAUSE\n')
        f.close()
        os.path.abspath = self.model_ws
        sp.call(batch_file_name, cwd=self.model_ws, stdout = report)
        os.remove(batch_file_name)
        return

    def run_model2(self, silent=True, pause=False, report=False):
        """
        This method will run the model using subprocess.Popen.

        Parameters
        ----------
        silent : boolean
            Echo run information to screen (default is True).
        pause : boolean, optional
            Pause upon completion (the default is False).
        report : string, optional
            Name of file to store stdout. (default is None).

        Returns
        -------
        (success, buff)
        success : boolean
        buff : list of lines of stdout

        """
        success = False
        buff = []
        proc = sp.Popen([self.exe_name,self.namefile],
                        stdout=sp.PIPE, cwd=self.model_ws)
        while True:
          line = proc.stdout.readline()
          if line != '':
            if 'normal termination of simulation' in line.lower():
                success = True
            #c = line.split('\r')
            c = line.rstrip('\r\n')
            if not silent:
                #print c[0]
                print c
            if report == True:
                buff.append( c[0] )
          else:
            break
        if pause == True:
            raw_input('Press Enter to continue...')
        return ( [success,buff] )

    def run_model3(self):
        """
        Minimal form of run model.  No input parameters and no return
        parameters.  Runs model using subprocess.Popen.
        """
        import subprocess
        a = subprocess.Popen([self.exe_name,self.namefile],
                             stdout=subprocess.PIPE, cwd=self.model_ws)
        b = a.communicate()
        c = b[0].split('\r')
        for cc in c:
            print cc

    def write_input(self, SelPackList=False):
        if self.verbose:
            print self # Same as calling self.__repr__()
            print 'Writing packages:'
        if SelPackList == False:
            for p in self.packagelist:
                p.write_file()
                if self.verbose:
                    print p.__repr__()
        else:
#            for i,p in enumerate(self.packagelist):
#                for pon in SelPackList:
            for pon in SelPackList:
                for i,p in enumerate(self.packagelist):
                    if pon in p.name:
                        print 'writing Package: ',p.name
                        p.write_file()
                        if self.verbose:
                            print p.__repr__()
                        break
        #--write name file
        self.write_name_file()

    def write_name_file(self):
        '''Every Package needs its own writenamefile function'''
        raise Exception, 'IMPLEMENTATION ERROR: writenamefile must be overloaded'

    def get_name(self):
        return self.__name

    def set_name(self, value):
        self.__name = value
        self.namefile = self.__name + '.' + self.namefile_ext
        for p in self.packagelist:
            for i in range(len(p.extension)):
                p.file_name[i] = self.__name + '.' + p.extension[i]
    name = property(get_name, set_name)

class Package(object):
    '''
    General Package class
      allowDuplicates allows more than one package of the same class to be added.
      This is needed for mfaddoutsidefile if used for more than one file.
    '''
    def __init__(self, parent, extension='glo', name='GLOBAL', unit_number=1, extra='',
                 allowDuplicates=False):
        self.parent = parent # To be able to access the parent modflow object's attributes
        if (not isinstance(extension, list)):
            extension = [extension]
        self.extension = []
        self.file_name = []
        for e in extension:
            self.extension = self.extension + [e]
            self.file_name = self.file_name + [self.parent.name + '.' + e]
            self.fn_path = os.path.join(self.parent.model_ws,self.file_name[0])
        if (not isinstance(name, list)):
            name = [name]
        self.name = name
        if (not isinstance(unit_number, list)):
            unit_number = [unit_number]
        self.unit_number = unit_number
        if (not isinstance(extra, list)):
            self.extra = len(self.unit_number) * [extra]
        else:
            self.extra = extra
        self.url = 'index.html'
        self.allowDuplicates = allowDuplicates
    def __repr__( self ):
        s = self.__doc__
        exclude_attributes = ['extension', 'heading', 'name', 'parent', 'url']
        for attr, value in sorted(self.__dict__.iteritems()):
            if not (attr in exclude_attributes):
                if (isinstance(value, list)):
                    if (len(value) == 1):
                        #s = s + ' %s = %s (list)\n' % (attr, str(value[0]))
                        s = s + ' {0:s} = {1:s}\n'.format(attr,str(value[0]))
                    else:
                        #s = s + ' %s (list, items = %d)\n' % (attr, len(value))
                        s = s + ' {0:s} (list, items = {1:d}\n'.format(attr,len(value))
                elif (isinstance(value, np.ndarray)):
                    #s = s + ' %s (array, shape = %s)\n' % (attr, value.shape.__str__()[1:-1] )
                    s = s + ' {0:s} (array, shape = {1:s}\n'.fomrat(attr,value.shape__str__()[1:-1])
                else:
                    #s = s + ' %s = %s (%s)\n' % (attr, str(value), str(type(value))[7:-2])
                    s = s + ' {0:s} = {1:s} ({2:s}\n'.format(attr,str(value),str(type(value))[7:-2])
        return s

    def assign_layer_row_column_data(self, layer_row_column_data, ncols, zerobase=True):
        if (layer_row_column_data is not None):
            new_layer_row_column_data = []
            mxact = 0
            for a in layer_row_column_data:
                a = np.atleast_2d(a)
                nr, nc = a.shape
                assert nc == ncols, 'layer_row_column_Q must have {0:1d} columns'.format(ncols)+'\nentry: '+str(a.shape)
                mxact = max(mxact, nr)
                if zerobase:
                    new_layer_row_column_data.append(a)
                else:
                    print 'Deprecation Warning: One-based indexing will be deprecated in future FloPy versions. Use Zero-based indexing'
                    a[:3,0] -= 1  # one-base input data, subtract 1 from layers, rows, columns
                    new_layer_row_column_data.append(a)
            return mxact, new_layer_row_column_data
        return

    def webdoc(self):
        if self.parent.version == 'mf2k':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow2000/Guide/' + self.url)
        elif self.parent.version == 'mf2005':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow2005/Guide/' + self.url)
        elif self.parent.version == 'ModflowNwt':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow_nwt/Guide/' + self.url)

    def write_file(self):
        '''Every Package needs its own write_file function'''
        print 'IMPLEMENTATION ERROR: write_file must be overloaded'

    def write_layer_row_column_data(self, f, layer_row_column_data):
        for n in range(self.parent.get_package('DIS').nper):
            if (n < len(layer_row_column_data)):
                a = layer_row_column_data[n]
                itmp = a.shape[0]
                #f.write('%10i%10i\n' % (itmp, self.np))
                f.write(' {0:9d} {1:9d}\n'.format(itmp,self.np))
                for b in a:
                    #f.write('%9i %9i %9i' % (b[0], b[1], b[2]) )
                    f.write(' {0:9d} {1:9d} {2:9d}\n'.format(b[0]+1,b[1],b[2]))  # write out layer+1
                    for c in b[3:]:
                        #f.write(' %13.6e' % c)
                        f.write(' {12.6}'.format(c))
                    f.write('\n')
            else:
                itmp = -1
                #f.write('%10i%10i\n' % (itmp, self.np))
                f.write(' {0:9d} {1:9d}\n'.format(itmp,self.np))
    def load(self):
        """
        The load method has not been implemented for this package.

        """
        print 'This package needs a load method()'
<commit_msg>zero base in assign array<commit_after>import numpy as np
import sys
import os
import subprocess as sp
import webbrowser as wb

# Global variables
iconst = 1 # Multiplier for individual array elements in integer and real arrays read by MODFLOW's U2DREL, U1DREL and U2DINT.
iprn = -1 # Printout flag. If >= 0 then array values read are printed in listing file.

class BaseModel(object):
    """
    MODFLOW based models base class
    """

    def __init__(self, modelname = 'modflowtest', namefile_ext = 'nam',
                 exe_name = 'mf2k.exe', model_ws = None):
        self.__name = modelname
        self.namefile_ext = namefile_ext
        self.namefile = self.__name + '.' + self.namefile_ext
        self.packagelist = []
        self.heading = ''
        self.exe_name = exe_name
        self.external_extension = 'ref'
        if model_ws is None: model_ws = os.getcwd()
        if not os.path.exists(model_ws):
            try:
                os.makedirs(model_ws)
            except:
                #print '\n%s not valid, workspace-folder was changed to %s\n' % (model_ws, os.getcwd())
                print '\n{0:s} not valid, workspace-folder was changed to {1:s}\n'.format(model_ws, os.getcwd())
                model_ws = os.getcwd()
        self.model_ws= model_ws
        self.cl_params = ''

    def set_exename(self, exe_name):
        self.exe_name = exe_name
        return

    def add_package(self, p):
        """
        Add a flopy package object to this model.
        """
        for pp in (self.packagelist):
            if pp.allowDuplicates:
                continue
            elif (isinstance(p, type(pp))):
                print '****Warning -- two packages of the same type: ',type(p),type(pp)
                print 'replacing existing Package...'
                pp = p
                return
        self.packagelist.append( p )

    def remove_package(self, pname):
        """
        Remove a package from this model.
        """
        for i,pp in enumerate(self.packagelist):
            if pname in pp.name:
                print 'removing Package: ',pp.name
                self.packagelist.pop(i)
                return
        raise StopIteration , 'Package name '+pname+' not found in Package list'

    def build_array_name(self,num,prefix):
       return self.external_path+prefix+'_'+str(num)+'.'+self.external_extension

    def assign_external(self,num,prefix):
        fname = self.build_array_name(num,prefix)
        unit = (self.next_ext_unit())
        self.external_fnames.append(fname)
        self.external_units.append(unit)
        self.external_binflag.append(False)
        return fname,unit

    def add_external(self, fname, unit, binflag=False):
        """
        Supports SWR usage and non-loaded existing external arrays
        """
        self.external_fnames.append(fname)
        self.external_units.append(unit)
        self.external_binflag.append(binflag)
        return

    def remove_external(self,fname=None,unit=None):
        if fname is not None:
            for i,e in enumerate(self.external_fnames):
                if fname in e:
                    self.external_fnames.pop(i)
                    self.external_units.pop(i)
                    self.external_binflag.pop(i)
        elif unit is not None:
            for i,u in enumerate(self.external_units):
                if u == unit:
                    self.external_fnames.pop(i)
                    self.external_units.pop(i)
                    self.external_binflag.pop(i)
        else:
            raise Exception,' either fname or unit must be passed to remove_external()'
        return

    def get_name_file_entries(self):
        s = ''
        for p in self.packagelist:
            for i in range(len(p.name)):
                #s = s + ('%s %3i %s %s\n' % (p.name[i], p.unit_number[i],
                #                             p.file_name[i],p.extra[i]))
                s = s + ('{0:s} {1:3d} {2:s} {3:s}\n'.format(p.name[i], p.unit_number[i],
                                             p.file_name[i],p.extra[i]))
        return s

    def get_package(self, name):
        for pp in (self.packagelist):
            if (pp.name[0].upper() == name.upper()):
                return pp
        return None

    def get_package_list(self):
        val = []
        for pp in (self.packagelist):
            val.append(pp.name[0].upper())
        return val

    def change_model_ws(self,new_pth=None):
        if new_pth is None:
            new_pth = os.getcwd()
        if not os.path.exists(new_pth):
            try:
                os.makedirs(new_pth)
            except:
                #print '\n%s not valid, workspace-folder was changed to %s\n' % (new_pth, os.getcwd())
                print '\n{0:s} not valid, workspace-folder was changed to {1:s}\n'.format(new_pth, os.getcwd())
                new_pth = os.getcwd()
        #--reset the model workspace
        self.model_ws = new_pth
        #--reset the paths for each package
        for pp in (self.packagelist):
            pp.fn_path = os.path.join(self.model_ws,pp.file_name[0])
        return None

    def run_model(self, pause=True, report=None):
        """
        Run the model.  This method will create and run a Windows batch file.

        Parameters
        ----------
        pause : boolean, optional
            Add a pause to the end of the batch file.
            (the default is False)
        report : string, optional
            Name of file to store stdout. (default is None).
        """
        batch_file_name = os.path.join(self.model_ws, 'run.bat')
        #error_message = ('Model executable %s not found!' % self.exe_name)
        error_message = 'Model executable {0:s} not found!'.format(self.exe_name)
        assert os.path.exists(self.exe_name), error_message

        error_message = ('Name file %s not found!' % self.namefile)
        fn_path = os.path.join(self.model_ws, self.namefile)
        assert os.path.exists(fn_path), error_message

        # Create a batch file to call code so that window remains open in case of error messages
        f = open(batch_file_name, 'w')
        #f.write('@ECHO Calling %s with %s\n' % (self.exe_name, self.namefile))
        f.write('@ECHO Calling {0:s} with {1:s}\n'.format(self.exe_name, self.namefile))
        #f.write('%s %s %s\n' % (self.exe_name, self.namefile, self.cl_params))
        f.write('{0:s} {1:s} {2:s}\n'.format(self.exe_name, self.namefile, self.cl_params))
        if (pause):
           f.write('@PAUSE\n')
        f.close()
        os.path.abspath = self.model_ws
        sp.call(batch_file_name, cwd=self.model_ws, stdout = report)
        os.remove(batch_file_name)
        return

    def run_model2(self, silent=True, pause=False, report=False):
        """
        This method will run the model using subprocess.Popen.

        Parameters
        ----------
        silent : boolean
            Echo run information to screen (default is True).
        pause : boolean, optional
            Pause upon completion (the default is False).
        report : string, optional
            Name of file to store stdout. (default is None).

        Returns
        -------
        (success, buff)
        success : boolean
        buff : list of lines of stdout

        """
        success = False
        buff = []
        proc = sp.Popen([self.exe_name,self.namefile],
                        stdout=sp.PIPE, cwd=self.model_ws)
        while True:
          line = proc.stdout.readline()
          if line != '':
            if 'normal termination of simulation' in line.lower():
                success = True
            #c = line.split('\r')
            c = line.rstrip('\r\n')
            if not silent:
                #print c[0]
                print c
            if report == True:
                buff.append( c[0] )
          else:
            break
        if pause == True:
            raw_input('Press Enter to continue...')
        return ( [success,buff] )

    def run_model3(self):
        """
        Minimal form of run model.  No input parameters and no return
        parameters.  Runs model using subprocess.Popen.
        """
        import subprocess
        a = subprocess.Popen([self.exe_name,self.namefile],
                             stdout=subprocess.PIPE, cwd=self.model_ws)
        b = a.communicate()
        c = b[0].split('\r')
        for cc in c:
            print cc

    def write_input(self, SelPackList=False):
        if self.verbose:
            print self # Same as calling self.__repr__()
            print 'Writing packages:'
        if SelPackList == False:
            for p in self.packagelist:
                p.write_file()
                if self.verbose:
                    print p.__repr__()
        else:
#            for i,p in enumerate(self.packagelist):
#                for pon in SelPackList:
            for pon in SelPackList:
                for i,p in enumerate(self.packagelist):
                    if pon in p.name:
                        print 'writing Package: ',p.name
                        p.write_file()
                        if self.verbose:
                            print p.__repr__()
                        break
        #--write name file
        self.write_name_file()

    def write_name_file(self):
        '''Every Package needs its own writenamefile function'''
        raise Exception, 'IMPLEMENTATION ERROR: writenamefile must be overloaded'

    def get_name(self):
        return self.__name

    def set_name(self, value):
        self.__name = value
        self.namefile = self.__name + '.' + self.namefile_ext
        for p in self.packagelist:
            for i in range(len(p.extension)):
                p.file_name[i] = self.__name + '.' + p.extension[i]
    name = property(get_name, set_name)

class Package(object):
    '''
    General Package class
      allowDuplicates allows more than one package of the same class to be added.
      This is needed for mfaddoutsidefile if used for more than one file.
    '''
    def __init__(self, parent, extension='glo', name='GLOBAL', unit_number=1, extra='',
                 allowDuplicates=False):
        self.parent = parent # To be able to access the parent modflow object's attributes
        if (not isinstance(extension, list)):
            extension = [extension]
        self.extension = []
        self.file_name = []
        for e in extension:
            self.extension = self.extension + [e]
            self.file_name = self.file_name + [self.parent.name + '.' + e]
            self.fn_path = os.path.join(self.parent.model_ws,self.file_name[0])
        if (not isinstance(name, list)):
            name = [name]
        self.name = name
        if (not isinstance(unit_number, list)):
            unit_number = [unit_number]
        self.unit_number = unit_number
        if (not isinstance(extra, list)):
            self.extra = len(self.unit_number) * [extra]
        else:
            self.extra = extra
        self.url = 'index.html'
        self.allowDuplicates = allowDuplicates
    def __repr__( self ):
        s = self.__doc__
        exclude_attributes = ['extension', 'heading', 'name', 'parent', 'url']
        for attr, value in sorted(self.__dict__.iteritems()):
            if not (attr in exclude_attributes):
                if (isinstance(value, list)):
                    if (len(value) == 1):
                        #s = s + ' %s = %s (list)\n' % (attr, str(value[0]))
                        s = s + ' {0:s} = {1:s}\n'.format(attr,str(value[0]))
                    else:
                        #s = s + ' %s (list, items = %d)\n' % (attr, len(value))
                        s = s + ' {0:s} (list, items = {1:d}\n'.format(attr,len(value))
                elif (isinstance(value, np.ndarray)):
                    #s = s + ' %s (array, shape = %s)\n' % (attr, value.shape.__str__()[1:-1] )
                    s = s + ' {0:s} (array, shape = {1:s}\n'.fomrat(attr,value.shape__str__()[1:-1])
                else:
                    #s = s + ' %s = %s (%s)\n' % (attr, str(value), str(type(value))[7:-2])
                    s = s + ' {0:s} = {1:s} ({2:s}\n'.format(attr,str(value),str(type(value))[7:-2])
        return s

    def assign_layer_row_column_data(self, layer_row_column_data, ncols, zerobase=True):
        if (layer_row_column_data is not None):
            new_layer_row_column_data = []
            mxact = 0
            for a in layer_row_column_data:
                a = np.atleast_2d(a)
                nr, nc = a.shape
                assert nc == ncols, 'layer_row_column_Q must have {0:1d} columns'.format(ncols)+'\nentry: '+str(a.shape)
                mxact = max(mxact, nr)
                if zerobase:
                    new_layer_row_column_data.append(a)
                else:
                    print 'Deprecation Warning: One-based indexing will be deprecated in future FloPy versions. Use Zero-based indexing'
                    a[:,:3] -= 1  # one-base input data, subtract 1 from layers, rows, columns
                    new_layer_row_column_data.append(a)
            return mxact, new_layer_row_column_data
        return

    def webdoc(self):
        if self.parent.version == 'mf2k':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow2000/Guide/' + self.url)
        elif self.parent.version == 'mf2005':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow2005/Guide/' + self.url)
        elif self.parent.version == 'ModflowNwt':
            wb.open('http://water.usgs.gov/nrp/gwsoftware/modflow_nwt/Guide/' + self.url)

    def write_file(self):
        '''Every Package needs its own write_file function'''
        print 'IMPLEMENTATION ERROR: write_file must be overloaded'

    def write_layer_row_column_data(self, f, layer_row_column_data):
        for n in range(self.parent.get_package('DIS').nper):
            if (n < len(layer_row_column_data)):
                a = layer_row_column_data[n]
                itmp = a.shape[0]
                #f.write('%10i%10i\n' % (itmp, self.np))
                f.write(' {0:9d} {1:9d}\n'.format(itmp,self.np))
                for b in a:
                    #f.write('%9i %9i %9i' % (b[0], b[1], b[2]) )
                    f.write(' {0:9d} {1:9d} {2:9d}\n'.format(b[0]+1,b[1],b[2]))  # write out layer+1
                    for c in b[3:]:
                        #f.write(' %13.6e' % c)
                        f.write(' {12.6}'.format(c))
                    f.write('\n')
            else:
                itmp = -1
                #f.write('%10i%10i\n' % (itmp, self.np))
                f.write(' {0:9d} {1:9d}\n'.format(itmp,self.np))
    def load(self):
        """
        The load method has not been implemented for this package.

        """
        print 'This package needs a load method()'
<|endoftext|>


<commit_before>from argparse import ArgumentParser
from tempfile import TemporaryDirectory
from pathlib import Path
import resource
import psutil
import mmap
import sys

import numpy as np
from astropy.io import fits

# This bit of hackery ensures that we can see ccdproc from within
# the test suite
sys.path.append(str(Path().cwd()))
from ccdproc import combine

# Do not combine these into one statement. When all references are lost
# to a TemporaryDirectory the directory is automatically deleted. _TMPDIR
# creates a reference that will stick around.
_TMPDIR = TemporaryDirectory()
TMPPATH = Path(_TMPDIR.name)

ALLOWED_EXTENSIONS = {
    'fits': 'fits',
    'plain': 'txt'
}

def generate_fits_files(number, size=None):
    if size is None:
        use_size = [250, 250]
    else:
        int_size = int(size)
        use_size = [int_size, int_size]

    base_name = 'test-combine-{num:03d}.' + ALLOWED_EXTENSIONS['fits']

    for num in range(number):
        data = np.zeros(shape=use_size)
        hdu = fits.PrimaryHDU(data=data)
        hdu.header['bunit'] = 'adu'
        name = base_name.format(num=num)
        path = TMPPATH / name
        hdu.writeto(path, overwrite=True)

def generate_plain_files(number):
    for i in range(number):
        file = TMPPATH / ("{i:03d}.".format(i=i) + ALLOWED_EXTENSIONS['plain'])
        file.write_bytes(np.random.random(100))

def open_files_with_open(kind):
    """
    Open files with plain open.
    """
    # Ensure the file references persist until end of script. Not really
    # necessary, but convenient while debugging the script.
    global fds
    fds = []

    paths = TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])

    for p in paths:
        fds.append(p.open())

def open_files_as_mmap(kind):
    """
    Open files as mmaps.
    """
    # Ensure the file references persist until end of script. Not really
    # necessary, but convenient while debugging the script.
    global fds
    fds = []

    paths = TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])

    for p in paths:
        with p.open() as f:
            fds.append(mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY))

def open_files_ccdproc_combine_chunk(kind):
    """
    Open files indirectly as part of ccdproc.combine, ensuring that the
    task is broken into chunks.
    """
    global combo
    paths = sorted(list(TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])))
    # We want to force combine to break the task into chunks even
    # if the task really would fit in memory; it is in that case that
    # we end up with too many open files. We'll open one file, determine
    # the size of the data in bytes, and set the memory limit to that.
    # That will mean lots of chunks (however many files there are plus one),
    # but lots of chunks is fine.
    with fits.open(paths[0]) as hdulist:
        array_size = hdulist[0].data.nbytes

    combo = combine(paths, mem_limit=array_size)

def open_files_ccdproc_combine_nochunk(kind):
    """
    Open files indirectly as part of ccdproc.combine, ensuring that the
    task is not broken into chunks.
    """
    global combo
    paths = sorted(list(TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])))

    # We ensure there are no chunks by setting a memory limit large
    # enough to hold everything.
    with fits.open(paths[0]) as hdulist:
        array_size = hdulist[0].data.nbytes

    # Why 2x the number of files? To make absolutely sure we don't
    # end up chunking the job.
    array_size *= 2 * len(paths)
    combo = combine(paths)

ALLOWED_OPENERS = {
    'open': open_files_with_open,
    'mmap': open_files_as_mmap,
    'combine-chunk': open_files_ccdproc_combine_chunk,
    'combine-nochunk': open_files_ccdproc_combine_nochunk
}

def run_with_limit(n, kind='fits', size=None, overhead=6,
                   open_method='mmap'):
    """
    Try opening a bunch of files with a relatively low limit on the number
    of open files.

    Parameters
    ----------

    n : int
        Limit on number of open files in this function. The number of files
        to create is calculated from this to be just below the maximum number
        of files controlled by this function that can be opened.

    kind : one of 'fits', 'plain', optional
        The type of file to generate. The plain files are intended mainly for
        testing this script, while the FITS files are for testing
        ccdproc.combine.

    size : int, optional
        Size of file to create. If the kind is 'plain; this is the size
        of the file, in bytes. If the kind is 'fits', this is the size
        of one side of the image (the image is always square).

    overhead : int, optional
        Number of open files to assume the OS is using for this process. The
        default value is chosen so that this succeeds on MacOS or Linux.
        Setting it to a value lower than default should cause a SystemExit
        exception to be raised because of too many open files. This is meant
        for testing that this script is actually testing something.

    Notes
    -----

    .. warning::

        You should run this in a subprocess. Running as part of a larger python
        process will lower the limit on the number of open files for that
        **entire python process** which will almost certainly lead to nasty
        side effects.
    """
    # Do a little input validation
    if n <= 0:
        raise ValueError("Argument 'n' must be a positive integer")

    if kind not in ALLOWED_EXTENSIONS.keys():
        raise ValueError("Argument 'kind' must be one of "
                         "{}".format(ALLOWED_EXTENSIONS.keys()))

    # Set the limit on the number of open files to n. The try/except
    # is the catch the case where this change would *increase*, rather than
    # decrease, the limit. That apparently can only be done by a superuser.
    try:
        resource.setrlimit(resource.RLIMIT_NOFILE, (n, n))
    except ValueError as e:
        if 'not allowed to raise maximum limit' not in str(e):
            raise
        max_n_this_process = resource.getrlimit(resource.RLIMIT_NOFILE)
        raise ValueError('Maximum number of open '
                         'files is {}'.format(max_n_this_process))

    # The "-1" is to leave a little wiggle room. overhead is based on the
    # the number of open files that a process running on linux has open.
    # These typically include stdin and stout, and apparently others.
    n_files = n - 1 - overhead

    proc = psutil.Process()

    print('Process ID is: ', proc.pid, flush=True)
    print("Making {} files".format(n_files))
    if kind == 'plain':
        generate_plain_files(n_files)
    elif kind == 'fits':
        generate_fits_files(n_files, size=size)

    # Print number of open files before we try opening anything for debugging
    # purposes.
    print("Before opening, files open is {}".format(len(proc.open_files())),
          flush=True)
    print("    Note well: this number is different than what lsof reports.")

    try:
        ALLOWED_OPENERS[open_method](kind)
            # fds.append(p.open())
    except OSError as e:
        # Capture the error and re-raise as a SystemExit because this is
        # run in a subprocess. This ensures that the original error message
        # is reported back to the calling process; we add on the number of
        # open files.
        raise SystemExit(str(e) + '; number of open files: ' +
                         '{}, with target {}'.format(len(proc.open_files()),
                                                     n_files))
    else:
        print('Opens succeeded, files currently open:',
              len(proc.open_files()),
              flush=True)

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('number', type=int,
                        help='Limit on number of open files.')
    parser.add_argument('--kind', action='store', default='plain',
                        choices=ALLOWED_EXTENSIONS.keys(),
                        help='Kind of file to generate for test; '
                             'default is plain')
    parser.add_argument('--overhead', type=int, action='store',
                        help='Number of files to assume the OS is using.',
                        default=6)
    parser.add_argument('--open-by', action='store', default='mmap',
                        choices=ALLOWED_OPENERS.keys(),
                        help='How to open the files. Default is mmap')
    parser.add_argument('--size', type=int, action='store',
                        help='Size of one side of image to create. '
                             'All images are square, so only give '
                             'a single number for the size.')
    args = parser.parse_args()
    run_with_limit(args.number, kind=args.kind, overhead=args.overhead,
                   open_method=args.open_by, size=args.size)
<commit_msg>Add option to change garbage collection frequency<commit_after>from argparse import ArgumentParser
from tempfile import TemporaryDirectory
from pathlib import Path
import resource
import psutil
import mmap
import sys
import gc

import numpy as np
from astropy.io import fits

# This bit of hackery ensures that we can see ccdproc from within
# the test suite
sys.path.append(str(Path().cwd()))
from ccdproc import combine

# Do not combine these into one statement. When all references are lost
# to a TemporaryDirectory the directory is automatically deleted. _TMPDIR
# creates a reference that will stick around.
_TMPDIR = TemporaryDirectory()
TMPPATH = Path(_TMPDIR.name)

ALLOWED_EXTENSIONS = {
    'fits': 'fits',
    'plain': 'txt'
}

def generate_fits_files(number, size=None):
    if size is None:
        use_size = [250, 250]
    else:
        int_size = int(size)
        use_size = [int_size, int_size]

    base_name = 'test-combine-{num:03d}.' + ALLOWED_EXTENSIONS['fits']

    for num in range(number):
        data = np.zeros(shape=use_size)
        hdu = fits.PrimaryHDU(data=data)
        hdu.header['bunit'] = 'adu'
        name = base_name.format(num=num)
        path = TMPPATH / name
        hdu.writeto(path, overwrite=True)

def generate_plain_files(number):
    for i in range(number):
        file = TMPPATH / ("{i:03d}.".format(i=i) + ALLOWED_EXTENSIONS['plain'])
        file.write_bytes(np.random.random(100))

def open_files_with_open(kind):
    """
    Open files with plain open.
    """
    # Ensure the file references persist until end of script. Not really
    # necessary, but convenient while debugging the script.
    global fds
    fds = []

    paths = TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])

    for p in paths:
        fds.append(p.open())

def open_files_as_mmap(kind):
    """
    Open files as mmaps.
    """
    # Ensure the file references persist until end of script. Not really
    # necessary, but convenient while debugging the script.
    global fds
    fds = []

    paths = TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])

    for p in paths:
        with p.open() as f:
            fds.append(mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY))

def open_files_ccdproc_combine_chunk(kind):
    """
    Open files indirectly as part of ccdproc.combine, ensuring that the
    task is broken into chunks.
    """
    global combo
    paths = sorted(list(TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])))
    # We want to force combine to break the task into chunks even
    # if the task really would fit in memory; it is in that case that
    # we end up with too many open files. We'll open one file, determine
    # the size of the data in bytes, and set the memory limit to that.
    # That will mean lots of chunks (however many files there are plus one),
    # but lots of chunks is fine.
    with fits.open(paths[0]) as hdulist:
        array_size = hdulist[0].data.nbytes

    combo = combine(paths, mem_limit=array_size)

def open_files_ccdproc_combine_nochunk(kind):
    """
    Open files indirectly as part of ccdproc.combine, ensuring that the
    task is not broken into chunks.
    """
    global combo
    paths = sorted(list(TMPPATH.glob('**/*.' + ALLOWED_EXTENSIONS[kind])))

    # We ensure there are no chunks by setting a memory limit large
    # enough to hold everything.
    with fits.open(paths[0]) as hdulist:
        array_size = hdulist[0].data.nbytes

    # Why 2x the number of files? To make absolutely sure we don't
    # end up chunking the job.
    array_size *= 2 * len(paths)
    combo = combine(paths)

ALLOWED_OPENERS = {
    'open': open_files_with_open,
    'mmap': open_files_as_mmap,
    'combine-chunk': open_files_ccdproc_combine_chunk,
    'combine-nochunk': open_files_ccdproc_combine_nochunk
}

def run_with_limit(n, kind='fits', size=None, overhead=6,
                   open_method='mmap'):
    """
    Try opening a bunch of files with a relatively low limit on the number
    of open files.

    Parameters
    ----------

    n : int
        Limit on number of open files in this function. The number of files
        to create is calculated from this to be just below the maximum number
        of files controlled by this function that can be opened.

    kind : one of 'fits', 'plain', optional
        The type of file to generate. The plain files are intended mainly for
        testing this script, while the FITS files are for testing
        ccdproc.combine.

    size : int, optional
        Size of file to create. If the kind is 'plain; this is the size
        of the file, in bytes. If the kind is 'fits', this is the size
        of one side of the image (the image is always square).

    overhead : int, optional
        Number of open files to assume the OS is using for this process. The
        default value is chosen so that this succeeds on MacOS or Linux.
        Setting it to a value lower than default should cause a SystemExit
        exception to be raised because of too many open files. This is meant
        for testing that this script is actually testing something.

    Notes
    -----

    .. warning::

        You should run this in a subprocess. Running as part of a larger python
        process will lower the limit on the number of open files for that
        **entire python process** which will almost certainly lead to nasty
        side effects.
    """
    # Do a little input validation
    if n <= 0:
        raise ValueError("Argument 'n' must be a positive integer")

    if kind not in ALLOWED_EXTENSIONS.keys():
        raise ValueError("Argument 'kind' must be one of "
                         "{}".format(ALLOWED_EXTENSIONS.keys()))

    # Set the limit on the number of open files to n. The try/except
    # is the catch the case where this change would *increase*, rather than
    # decrease, the limit. That apparently can only be done by a superuser.
    try:
        resource.setrlimit(resource.RLIMIT_NOFILE, (n, n))
    except ValueError as e:
        if 'not allowed to raise maximum limit' not in str(e):
            raise
        max_n_this_process = resource.getrlimit(resource.RLIMIT_NOFILE)
        raise ValueError('Maximum number of open '
                         'files is {}'.format(max_n_this_process))

    # The "-1" is to leave a little wiggle room. overhead is based on the
    # the number of open files that a process running on linux has open.
    # These typically include stdin and stout, and apparently others.
    n_files = n - 1 - overhead

    proc = psutil.Process()

    print('Process ID is: ', proc.pid, flush=True)
    print("Making {} files".format(n_files))
    if kind == 'plain':
        generate_plain_files(n_files)
    elif kind == 'fits':
        generate_fits_files(n_files, size=size)

    # Print number of open files before we try opening anything for debugging
    # purposes.
    print("Before opening, files open is {}".format(len(proc.open_files())),
          flush=True)
    print("    Note well: this number is different than what lsof reports.")

    try:
        ALLOWED_OPENERS[open_method](kind)
            # fds.append(p.open())
    except OSError as e:
        # Capture the error and re-raise as a SystemExit because this is
        # run in a subprocess. This ensures that the original error message
        # is reported back to the calling process; we add on the number of
        # open files.
        raise SystemExit(str(e) + '; number of open files: ' +
                         '{}, with target {}'.format(len(proc.open_files()),
                                                     n_files))
    else:
        print('Opens succeeded, files currently open:',
              len(proc.open_files()),
              flush=True)

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('number', type=int,
                        help='Limit on number of open files.')
    parser.add_argument('--kind', action='store', default='plain',
                        choices=ALLOWED_EXTENSIONS.keys(),
                        help='Kind of file to generate for test; '
                             'default is plain')
    parser.add_argument('--overhead', type=int, action='store',
                        help='Number of files to assume the OS is using.',
                        default=6)
    parser.add_argument('--open-by', action='store', default='mmap',
                        choices=ALLOWED_OPENERS.keys(),
                        help='How to open the files. Default is mmap')
    parser.add_argument('--size', type=int, action='store',
                        help='Size of one side of image to create. '
                             'All images are square, so only give '
                             'a single number for the size.')
    parser.add_argument('--frequent-gc', action='store_true',
                        help='If set, perform garbage collection '
                             'much more frequently than the default.')
    args = parser.parse_args()
    if args.frequent_gc:
        gc.set_threshold(10, 10, 10)
    print("Garbage collection thresholds: ", gc.get_threshold())
    run_with_limit(args.number, kind=args.kind, overhead=args.overhead,
                   open_method=args.open_by, size=args.size)
<|endoftext|>


<commit_before># Copyright (C) 2017 FireEye, Inc. All Rights Reserved.

import re
import tabulate
from collections import OrderedDict

ONE_MB = 1024 * 1024
STACK_MEM_NAME = "[stack]"

MAX_STRING_LENGTH = 2048

def makeEmulator(vw):
    """
    create an emulator using consistent settings.
    """
    emu = vw.getEmulator(logwrite=True)
    removeStackMemory(emu)
    emu.initStackMemory(stacksize=int(0.5 * ONE_MB))
    emu.setStackCounter(emu.getStackCounter() - int(0.25 * ONE_MB))
    emu.setEmuOpt('i386:reponce', False)  # do not short circuit rep prefix
    return emu

def removeStackMemory(emu):
    # TODO this is a hack while vivisect's initStackMemory() has a bug (see issue #27)
    # TODO does this bug still exist?
    memory_snap = emu.getMemorySnap()
    for i in xrange((len(memory_snap) - 1), -1, -1):
        (_, _, info, _) = memory_snap[i]
        if info[3] == STACK_MEM_NAME:
            del memory_snap[i]
            emu.setMemorySnap(memory_snap)
            emu.stack_map_base = None
            return
    raise Exception  # STACK_MEM_NAME not in memory map

def get_vivisect_meta_info(vw, selected_functions):
    info = OrderedDict()
    entry_points = vw.getEntryPoints()
    basename = None
    if entry_points:
        basename = vw.getFileByVa(entry_points[0])
    if basename:
        version = vw.getFileMeta(basename, 'Version')
        md5sum = vw.getFileMeta(basename, 'md5sum')
        baseva = hex(vw.getFileMeta(basename, 'imagebase'))
    else:
        version = "N/A"
        md5sum = "N/A"
        baseva = "N/A"

    info["Version"] = version
    info["MD5 Sum"] = md5sum
    info["Format"] = vw.getMeta("Format")
    info["Architecture"] = vw.getMeta("Architecture")
    info["Platform"] = vw.getMeta("Platform")
    disc, undisc = vw.getDiscoveredInfo()
    info["Percentage of discovered executable surface area"] = "%.1f%% (%s / %s)" % (disc * 100.0 / (disc + undisc), disc, disc + undisc)
    info["Base VA"] = baseva
    info["Entry point(s)"] = ", ".join(map(hex, entry_points))
    info["Number of imports"] = len(vw.getImports())
    info["Number of exports"] = len(vw.getExports())
    info["Number of functions"] = len(vw.getFunctions())
    if selected_functions:
        meta = []
        for fva in selected_functions:
            xrefs_to = len(vw.getXrefsTo(fva))
            num_args = len(vw.getFunctionArgs(fva))
            function_meta = vw.getFunctionMetaDict(fva)
            instr_count = function_meta["InstructionCount"]
            block_count = function_meta["BlockCount"]
            size = function_meta["Size"]
            meta.append((hex(fva), xrefs_to, num_args, size, block_count, instr_count))
        info["Selected functions' info"] = "\n%s" % tabulate.tabulate(meta, headers=["fva", "#xrefs", "#args", "size", "#blocks", "#instructions"])
    return info

def hex(i):
    return "0x%X" % (i)

FP_FILTER_PREFIXES = re.compile(r"^.?((p|P|0)?VA)|(0|P)?\\A|\[A|P\]A|@AA")  # remove string prefixes: pVA, VA, 0VA, etc.
FP_FILTER_SUFFIXES = re.compile(r".*([0-9A-G>]VA$|@AA)$")  # remove string suffixes: 0VA, AVA, >VA, etc.
FP_FILTER_CHARS = re.compile(r".*(AAA|BBB|CCC|DDD|EEE|FFF|PPP|UUU|ZZZ|@@@|;;;|&&&|\?\?\?|\|\|\||    ).*")
# alternatively: ".*([^0-9wW])\1{2}.*" to match any 3 consecutive chars (except numbers, ws, and others?)

def is_fp_string(s):
    """
    Return True if string matches a well-known FP pattern.
    :param s: input string
    """
    return FP_FILTER_CHARS.match(s)

def strip_string(s):
    """
    Return string stripped from false positive (FP) pre- or suffixes.
    :param s: input string
    :return: string stripped from FP pre- or suffixes
    """
    for reg in (FP_FILTER_PREFIXES, FP_FILTER_SUFFIXES):
        s = re.sub(reg, "", s)
    return s
<commit_msg>more FP filters<commit_after># Copyright (C) 2017 FireEye, Inc. All Rights Reserved.

import re
import tabulate
from collections import OrderedDict

ONE_MB = 1024 * 1024
STACK_MEM_NAME = "[stack]"

MAX_STRING_LENGTH = 2048

def makeEmulator(vw):
    """
    create an emulator using consistent settings.
    """
    emu = vw.getEmulator(logwrite=True)
    removeStackMemory(emu)
    emu.initStackMemory(stacksize=int(0.5 * ONE_MB))
    emu.setStackCounter(emu.getStackCounter() - int(0.25 * ONE_MB))
    emu.setEmuOpt('i386:reponce', False)  # do not short circuit rep prefix
    return emu

def removeStackMemory(emu):
    # TODO this is a hack while vivisect's initStackMemory() has a bug (see issue #27)
    # TODO does this bug still exist?
    memory_snap = emu.getMemorySnap()
    for i in xrange((len(memory_snap) - 1), -1, -1):
        (_, _, info, _) = memory_snap[i]
        if info[3] == STACK_MEM_NAME:
            del memory_snap[i]
            emu.setMemorySnap(memory_snap)
            emu.stack_map_base = None
            return
    raise Exception  # STACK_MEM_NAME not in memory map

def get_vivisect_meta_info(vw, selected_functions):
    info = OrderedDict()
    entry_points = vw.getEntryPoints()
    basename = None
    if entry_points:
        basename = vw.getFileByVa(entry_points[0])
    if basename:
        version = vw.getFileMeta(basename, 'Version')
        md5sum = vw.getFileMeta(basename, 'md5sum')
        baseva = hex(vw.getFileMeta(basename, 'imagebase'))
    else:
        version = "N/A"
        md5sum = "N/A"
        baseva = "N/A"

    info["Version"] = version
    info["MD5 Sum"] = md5sum
    info["Format"] = vw.getMeta("Format")
    info["Architecture"] = vw.getMeta("Architecture")
    info["Platform"] = vw.getMeta("Platform")
    disc, undisc = vw.getDiscoveredInfo()
    info["Percentage of discovered executable surface area"] = "%.1f%% (%s / %s)" % (disc * 100.0 / (disc + undisc), disc, disc + undisc)
    info["Base VA"] = baseva
    info["Entry point(s)"] = ", ".join(map(hex, entry_points))
    info["Number of imports"] = len(vw.getImports())
    info["Number of exports"] = len(vw.getExports())
    info["Number of functions"] = len(vw.getFunctions())
    if selected_functions:
        meta = []
        for fva in selected_functions:
            xrefs_to = len(vw.getXrefsTo(fva))
            num_args = len(vw.getFunctionArgs(fva))
            function_meta = vw.getFunctionMetaDict(fva)
            instr_count = function_meta["InstructionCount"]
            block_count = function_meta["BlockCount"]
            size = function_meta["Size"]
            meta.append((hex(fva), xrefs_to, num_args, size, block_count, instr_count))
        info["Selected functions' info"] = "\n%s" % tabulate.tabulate(meta, headers=["fva", "#xrefs", "#args", "size", "#blocks", "#instructions"])
    return info

def hex(i):
    return "0x%X" % (i)

FP_FILTER_PREFIXES = re.compile(r"^.?((p|P|0)?VA)|(0|P)?\\A|\[A|P\]A|@AA")  # remove string prefixes: pVA, VA, 0VA, etc.
FP_FILTER_SUFFIXES = re.compile(r".*([0-9A-G>]VA$|@AA|iiVV)$")  # remove string suffixes: 0VA, AVA, >VA, etc.
FP_FILTER_CHARS = re.compile(r".*(AAA|BBB|CCC|DDD|EEE|FFF|PPP|UUU|ZZZ|@@@|;;;|&&&|\?\?\?|\|\|\||    ).*")
# alternatively: ".*([^0-9wW])\1{2}.*" to match any 3 consecutive chars (except numbers, ws, and others?)
FP_FILTER_REP_CHARS = re.compile(r".*(.)\1{7}.*")  # any string containing the same char 8 or more consecutive times

def is_fp_string(s):
    """
    Return True if string matches a well-known FP pattern.
    :param s: input string
    """
    for reg in (FP_FILTER_CHARS, FP_FILTER_REP_CHARS):
        if reg.match(s):
            return True
    return False

def strip_string(s):
    """
    Return string stripped from false positive (FP) pre- or suffixes.
    :param s: input string
    :return: string stripped from FP pre- or suffixes
    """
    for reg in (FP_FILTER_PREFIXES, FP_FILTER_SUFFIXES):
        s = re.sub(reg, "", s)
    return s
<|endoftext|>


<commit_before>import argparse
import fnmatch
import json
import os
import re
import sys

parser = argparse.ArgumentParser(description='Find text between markdown JSON tags')
parser.add_argument('--extension', help='Extension of files to validate')
parser.add_argument('--directory', help='Path of directory to recursively go through')
args = parser.parse_args()

exit_status = 0

for root, dirs, files in os.walk(args.directory):
    for filename in fnmatch.filter(files, '*' + args.extension):
        with open(os.path.join(root, filename), "r") as validation_file:

            validation_data = validation_file.read()
            x = re.findall(r'{{< highlight json >}}(.*?){{< /highlight >}}',
                           validation_data, re.DOTALL)

            for json_block in x:
                try:
                    json.loads(json_block)
                except ValueError as exception:
                    print "In file " + validation_file.name + " the following JSON is invalid\n" + \
                    str(exception) + "\n" + json_block + "\n"
                    exit_status = 2
sys.exit(exit_status)
<commit_msg>Updates json validator script to use new code tags (#2547)<commit_after>import argparse
import fnmatch
import json
import os
import re
import sys

parser = argparse.ArgumentParser(description='Find text between markdown JSON tags')
parser.add_argument('--extension', help='Extension of files to validate')
parser.add_argument('--directory', help='Path of directory to recursively go through')
args = parser.parse_args()

exit_status = 0

for root, dirs, files in os.walk(args.directory):
    for filename in fnmatch.filter(files, '*' + args.extension):
        with open(os.path.join(root, filename), "r") as validation_file:

            validation_data = validation_file.read()
            x = re.findall(r'{{< code json >}}(.*?){{< /code >}}',
                           validation_data, re.DOTALL)

            for json_block in x:
                try:
                    json.loads(json_block)
                except ValueError as exception:
                    print "In file " + validation_file.name + " the following JSON is invalid\n" + \
                    str(exception) + "\n" + json_block + "\n"
                    exit_status = 2
sys.exit(exit_status)
<|endoftext|>


<commit_before>from spec import Spec, skip, ok_, eq_, raises, trap

from invoke.parser import Parser, Context, Argument, ParseError
from invoke.collection import Collection

class Parser_(Spec):
    def can_take_initial_context(self):
        c = Context()
        p = Parser(initial=c)
        eq_(p.initial, c)

    def can_take_initial_and_other_contexts(self):
        c1 = Context('foo')
        c2 = Context('bar')
        p = Parser(initial=Context(), contexts=[c1, c2])
        eq_(p.contexts['foo'], c1)
        eq_(p.contexts['bar'], c2)

    def can_take_just_other_contexts(self):
        c = Context('foo')
        p = Parser(contexts=[c])
        eq_(p.contexts['foo'], c)

    def can_take_just_contexts_as_non_keyword_arg(self):
        c = Context('foo')
        p = Parser([c])
        eq_(p.contexts['foo'], c)

    @raises(ValueError)
    def raises_ValueError_for_unnamed_Contexts_in_contexts(self):
        Parser(initial=Context(), contexts=[Context()])

    @raises(ValueError)
    def raises_error_for_context_name_clashes(self):
        Parser(contexts=(Context('foo'), Context('foo')))

    @raises(ValueError)
    def raises_error_for_context_alias_and_name_clashes(self):
        Parser((Context('foo', aliases=('bar',)), Context('bar')))

    @raises(ValueError)
    def raises_error_for_context_name_and_alias_clashes(self):
        # I.e. inverse of the above, which is a different code path.
        Parser((Context('foo'), Context('bar', aliases=('foo',))))

    def takes_ignore_unknown_kwarg(self):
        Parser(ignore_unknown=True)

    def ignore_unknown_defaults_to_False(self):
        eq_(Parser().ignore_unknown, False)

    class parse_argv:
        def parses_sys_argv_style_list_of_strings(self):
            "parses sys.argv-style list of strings"
            # Doesn't-blow-up tests FTL
            mytask = Context(name='mytask')
            mytask.add_arg('arg')
            p = Parser(contexts=[mytask])
            p.parse_argv(['mytask', '--arg', 'value'])

        def returns_only_contexts_mentioned(self):
            task1 = Context('mytask')
            task2 = Context('othertask')
            result = Parser((task1, task2)).parse_argv(['othertask'])
            eq_(len(result), 1)
            eq_(result[0].name, 'othertask')

        @raises(ParseError)
        def raises_error_if_unknown_contexts_found(self):
            Parser().parse_argv(['foo', 'bar'])

        def unparsed_does_not_share_state(self):
            r = Parser(ignore_unknown=True).parse_argv(['self'])
            eq_(r.unparsed, ['self'])
            r2 = Parser(ignore_unknown=True).parse_argv(['contained'])
            eq_(r.unparsed, ['self']) # NOT ['self', 'contained']
            eq_(r2.unparsed, ['contained']) # NOT ['self', 'contained']

        def ignore_unknown_returns_unparsed_argv_instead(self):
            r = Parser(ignore_unknown=True).parse_argv(['foo', 'bar', '--baz'])
            eq_(r.unparsed, ['foo', 'bar', '--baz'])

        def ignore_unknown_does_not_mutate_rest_of_argv(self):
            p = Parser([Context('ugh')], ignore_unknown=True)
            r = p.parse_argv(['ugh', 'what', '-nowai'])
            # NOT: ['what', '-n', '-w', '-a', '-i']
            eq_(r.unparsed, ['what', '-nowai'])

        def always_includes_initial_context_if_one_was_given(self):
            # Even if no core/initial flags were seen
            t1 = Context('t1')
            init = Context()
            result = Parser((t1,), initial=init).parse_argv(['t1'])
            eq_(result[0].name, None)
            eq_(result[1].name, 't1')

        def returned_contexts_are_in_order_given(self):
            t1, t2 = Context('t1'), Context('t2')
            r = Parser((t1, t2)).parse_argv(['t2', 't1'])
            eq_([x.name for x in r], ['t2', 't1'])

        def returned_context_member_arguments_contain_given_values(self):
            c = Context('mytask', args=(Argument('boolean', kind=bool),))
            result = Parser((c,)).parse_argv(['mytask', '--boolean'])
            eq_(result[0].args['boolean'].value, True)

        def arguments_which_take_values_get_defaults_overridden_correctly(self):
            args = (Argument('arg', kind=str), Argument('arg2', kind=int))
            c = Context('mytask', args=args)
            argv = ['mytask', '--arg', 'myval', '--arg2', '25']
            result = Parser((c,)).parse_argv(argv)
            eq_(result[0].args['arg'].value, 'myval')
            eq_(result[0].args['arg2'].value, 25)

        def returned_arguments_not_given_contain_default_values(self):
            # I.e. a Context with args A and B, invoked with no mention of B,
            # should result in B existing in the result, with its default value
            # intact, and not e.g. None, or the arg not existing.
            a = Argument('name', kind=str)
            b = Argument('age', default=7)
            c = Context('mytask', args=(a, b))
            result = Parser((c,)).parse_argv(['mytask', '--name', 'blah'])
            eq_(c.args['age'].value, 7)

        def returns_remainder(self):
            "returns -- style remainder string chunk"
            r = Parser((Context('foo'),)).parse_argv(['foo', '--', 'bar', 'biz'])
            eq_(r.remainder, "bar biz")

        def clones_initial_context(self):
            a = Argument('foo', kind=bool)
            eq_(a.value, None)
            c = Context(args=(a,))
            p = Parser(initial=c)
            assert p.initial is c
            r = p.parse_argv(['--foo'])
            assert p.initial is c
            c2 = r[0]
            assert c2 is not c
            a2 = c2.args['foo']
            assert a2 is not a
            eq_(a.value, None)
            eq_(a2.value, True)

        def clones_noninitial_contexts(self):
            a = Argument('foo')
            eq_(a.value, None)
            c = Context(name='mytask', args=(a,))
            p = Parser(contexts=(c,))
            assert p.contexts['mytask'] is c
            r = p.parse_argv(['mytask', '--foo', 'val'])
            assert p.contexts['mytask'] is c
            c2 = r[0]
            assert c2 is not c
            a2 = c2.args['foo']
            assert a2 is not a
            eq_(a.value, None)
            eq_(a2.value, 'val')

        class parsing_errors:
            def setup(self):
                self.p = Parser([Context(name='foo', args=[Argument('bar')])])

            @raises(ParseError)
            def missing_flag_values_raise_ParseError(self):
                self.p.parse_argv(['foo', '--bar'])

            def attaches_context_to_ParseErrors(self):
                try:
                    self.p.parse_argv(['foo', '--bar'])
                except ParseError as e:
                    assert e.context is not None

            def attached_context_is_None_outside_contexts(self):
                try:
                    Parser().parse_argv(['wat'])
                except ParseError as e:
                    assert e.context is None

        class positional_arguments:
            def _basic(self):
                arg = Argument('pos', positional=True)
                mytask = Context(name='mytask', args=[arg])
                return Parser(contexts=[mytask])

            def single_positional_arg(self):
                r = self._basic().parse_argv(['mytask', 'posval'])
                eq_(r[0].args['pos'].value, 'posval')

            @raises(ParseError)
            def omitted_positional_arg_raises_ParseError(self):
                self._basic().parse_argv(['mytask'])

            def positional_args_eat_otherwise_valid_context_names(self):
                mytask = Context('mytask', args=[
                    Argument('pos', positional=True),
                    Argument('nonpos', default='default')
                ])
                othertask = Context('lolwut')
                result = Parser([mytask]).parse_argv(['mytask', 'lolwut'])
                r = result[0]
                eq_(r.args['pos'].value, 'lolwut')
                eq_(r.args['nonpos'].value, 'default')
                eq_(len(result), 1) # Not 2

            def positional_args_can_still_be_given_as_flags(self):
                # AKA "positional args can come anywhere in the context"
                pos1 = Argument('pos1', positional=True)
                pos2 = Argument('pos2', positional=True)
                nonpos = Argument('nonpos', positional=False, default='lol')
                mytask = Context('mytask', args=[pos1, pos2, nonpos])
                eq_(mytask.positional_args, [pos1, pos2])
                r = Parser([mytask]).parse_argv([
                    'mytask',
                    '--nonpos', 'wut',
                    '--pos2', 'pos2val',
                    'pos1val',
                ])[0]
                eq_(r.args['pos1'].value, 'pos1val')
                eq_(r.args['pos2'].value, 'pos2val')
                eq_(r.args['nonpos'].value, 'wut')

        class equals_signs:
            def _compare(self, argname, invoke, value):
                c = Context('mytask', args=(Argument(argname, kind=str),))
                r = Parser((c,)).parse_argv(['mytask', invoke])
                eq_(r[0].args[argname].value, value)

            def handles_equals_style_long_flags(self):
                self._compare('foo', '--foo=bar', 'bar')

            def handles_equals_style_short_flags(self):
                self._compare('f', '-f=bar', 'bar')

            def does_not_require_escaping_equals_signs_in_value(self):
                self._compare('f', '-f=biz=baz', 'biz=baz')

        def handles_multiple_boolean_flags_per_context(self):
            c = Context('mytask', args=(
                Argument('foo', kind=bool), Argument('bar', kind=bool)
            ))
            r = Parser([c]).parse_argv(['mytask', '--foo', '--bar'])
            a = r[0].args
            eq_(a.foo.value, True)
            eq_(a.bar.value, True)

    class optional_arg_values:
        def setup(self):
            self.context = Context('mytask', args=(
                Argument('foo', optional=True, default='mydefault'),
            ))
            self.parser = Parser([self.context])

        def _expect(self, argv, expected):
            result = self.parser.parse_argv(['mytask'] + argv)
            eq_(result[0].args.foo.value, expected)

        def no_value_becomes_True(self):
            self._expect(['--foo'], True)

        def value_given_gets_preserved_normally(self):
            self._expect(['--foo', 'whatever'], 'whatever')

        def not_given_at_all_uses_default_value(self):
            self._expect([], 'mydefault')

        def ambiguity_with_unfilled_posargs(self):
            # mytask --foo value-not-posarg
            skip()

        def ambiguity_with_flaglike_value(self):
            # mytask --foo --lolwut
            skip()

        def ambiguity_with_task_name(self):
            # mytask --foo myothertask
            skip()

class ParseResult_(Spec):
    "ParseResult"
    def setup(self):
        self.context = Context('mytask',
            args=(Argument('foo', kind=str), Argument('bar')))
        argv = ['mytask', '--foo', 'foo-val', '--', 'my', 'remainder']
        self.result = Parser((self.context,)).parse_argv(argv)

    def acts_as_a_list_of_parsed_contexts(self):
        eq_(len(self.result), 1)
        eq_(self.result[0].name, 'mytask')

    def exhibits_remainder_attribute(self):
        eq_(self.result.remainder, 'my remainder')

    def to_dict_returns_parsed_contexts_and_args_as_nested_dicts(self):
        eq_(
            self.result.to_dict(),
            {'mytask': {'foo': 'foo-val', 'bar': None}}
        )
<commit_msg>Overhaul new tests to be more comprehensive<commit_after>from spec import Spec, skip, ok_, eq_, raises, trap

from invoke.parser import Parser, Context, Argument, ParseError
from invoke.collection import Collection

class Parser_(Spec):
    def can_take_initial_context(self):
        c = Context()
        p = Parser(initial=c)
        eq_(p.initial, c)

    def can_take_initial_and_other_contexts(self):
        c1 = Context('foo')
        c2 = Context('bar')
        p = Parser(initial=Context(), contexts=[c1, c2])
        eq_(p.contexts['foo'], c1)
        eq_(p.contexts['bar'], c2)

    def can_take_just_other_contexts(self):
        c = Context('foo')
        p = Parser(contexts=[c])
        eq_(p.contexts['foo'], c)

    def can_take_just_contexts_as_non_keyword_arg(self):
        c = Context('foo')
        p = Parser([c])
        eq_(p.contexts['foo'], c)

    @raises(ValueError)
    def raises_ValueError_for_unnamed_Contexts_in_contexts(self):
        Parser(initial=Context(), contexts=[Context()])

    @raises(ValueError)
    def raises_error_for_context_name_clashes(self):
        Parser(contexts=(Context('foo'), Context('foo')))

    @raises(ValueError)
    def raises_error_for_context_alias_and_name_clashes(self):
        Parser((Context('foo', aliases=('bar',)), Context('bar')))

    @raises(ValueError)
    def raises_error_for_context_name_and_alias_clashes(self):
        # I.e. inverse of the above, which is a different code path.
        Parser((Context('foo'), Context('bar', aliases=('foo',))))

    def takes_ignore_unknown_kwarg(self):
        Parser(ignore_unknown=True)

    def ignore_unknown_defaults_to_False(self):
        eq_(Parser().ignore_unknown, False)

    class parse_argv:
        def parses_sys_argv_style_list_of_strings(self):
            "parses sys.argv-style list of strings"
            # Doesn't-blow-up tests FTL
            mytask = Context(name='mytask')
            mytask.add_arg('arg')
            p = Parser(contexts=[mytask])
            p.parse_argv(['mytask', '--arg', 'value'])

        def returns_only_contexts_mentioned(self):
            task1 = Context('mytask')
            task2 = Context('othertask')
            result = Parser((task1, task2)).parse_argv(['othertask'])
            eq_(len(result), 1)
            eq_(result[0].name, 'othertask')

        @raises(ParseError)
        def raises_error_if_unknown_contexts_found(self):
            Parser().parse_argv(['foo', 'bar'])

        def unparsed_does_not_share_state(self):
            r = Parser(ignore_unknown=True).parse_argv(['self'])
            eq_(r.unparsed, ['self'])
            r2 = Parser(ignore_unknown=True).parse_argv(['contained'])
            eq_(r.unparsed, ['self']) # NOT ['self', 'contained']
            eq_(r2.unparsed, ['contained']) # NOT ['self', 'contained']

        def ignore_unknown_returns_unparsed_argv_instead(self):
            r = Parser(ignore_unknown=True).parse_argv(['foo', 'bar', '--baz'])
            eq_(r.unparsed, ['foo', 'bar', '--baz'])

        def ignore_unknown_does_not_mutate_rest_of_argv(self):
            p = Parser([Context('ugh')], ignore_unknown=True)
            r = p.parse_argv(['ugh', 'what', '-nowai'])
            # NOT: ['what', '-n', '-w', '-a', '-i']
            eq_(r.unparsed, ['what', '-nowai'])

        def always_includes_initial_context_if_one_was_given(self):
            # Even if no core/initial flags were seen
            t1 = Context('t1')
            init = Context()
            result = Parser((t1,), initial=init).parse_argv(['t1'])
            eq_(result[0].name, None)
            eq_(result[1].name, 't1')

        def returned_contexts_are_in_order_given(self):
            t1, t2 = Context('t1'), Context('t2')
            r = Parser((t1, t2)).parse_argv(['t2', 't1'])
            eq_([x.name for x in r], ['t2', 't1'])

        def returned_context_member_arguments_contain_given_values(self):
            c = Context('mytask', args=(Argument('boolean', kind=bool),))
            result = Parser((c,)).parse_argv(['mytask', '--boolean'])
            eq_(result[0].args['boolean'].value, True)

        def arguments_which_take_values_get_defaults_overridden_correctly(self):
            args = (Argument('arg', kind=str), Argument('arg2', kind=int))
            c = Context('mytask', args=args)
            argv = ['mytask', '--arg', 'myval', '--arg2', '25']
            result = Parser((c,)).parse_argv(argv)
            eq_(result[0].args['arg'].value, 'myval')
            eq_(result[0].args['arg2'].value, 25)

        def returned_arguments_not_given_contain_default_values(self):
            # I.e. a Context with args A and B, invoked with no mention of B,
            # should result in B existing in the result, with its default value
            # intact, and not e.g. None, or the arg not existing.
            a = Argument('name', kind=str)
            b = Argument('age', default=7)
            c = Context('mytask', args=(a, b))
            result = Parser((c,)).parse_argv(['mytask', '--name', 'blah'])
            eq_(c.args['age'].value, 7)

        def returns_remainder(self):
            "returns -- style remainder string chunk"
            r = Parser((Context('foo'),)).parse_argv(['foo', '--', 'bar', 'biz'])
            eq_(r.remainder, "bar biz")

        def clones_initial_context(self):
            a = Argument('foo', kind=bool)
            eq_(a.value, None)
            c = Context(args=(a,))
            p = Parser(initial=c)
            assert p.initial is c
            r = p.parse_argv(['--foo'])
            assert p.initial is c
            c2 = r[0]
            assert c2 is not c
            a2 = c2.args['foo']
            assert a2 is not a
            eq_(a.value, None)
            eq_(a2.value, True)

        def clones_noninitial_contexts(self):
            a = Argument('foo')
            eq_(a.value, None)
            c = Context(name='mytask', args=(a,))
            p = Parser(contexts=(c,))
            assert p.contexts['mytask'] is c
            r = p.parse_argv(['mytask', '--foo', 'val'])
            assert p.contexts['mytask'] is c
            c2 = r[0]
            assert c2 is not c
            a2 = c2.args['foo']
            assert a2 is not a
            eq_(a.value, None)
            eq_(a2.value, 'val')

        class parsing_errors:
            def setup(self):
                self.p = Parser([Context(name='foo', args=[Argument('bar')])])

            @raises(ParseError)
            def missing_flag_values_raise_ParseError(self):
                self.p.parse_argv(['foo', '--bar'])

            def attaches_context_to_ParseErrors(self):
                try:
                    self.p.parse_argv(['foo', '--bar'])
                except ParseError as e:
                    assert e.context is not None

            def attached_context_is_None_outside_contexts(self):
                try:
                    Parser().parse_argv(['wat'])
                except ParseError as e:
                    assert e.context is None

        class positional_arguments:
            def _basic(self):
                arg = Argument('pos', positional=True)
                mytask = Context(name='mytask', args=[arg])
                return Parser(contexts=[mytask])

            def single_positional_arg(self):
                r = self._basic().parse_argv(['mytask', 'posval'])
                eq_(r[0].args['pos'].value, 'posval')

            @raises(ParseError)
            def omitted_positional_arg_raises_ParseError(self):
                self._basic().parse_argv(['mytask'])

            def positional_args_eat_otherwise_valid_context_names(self):
                mytask = Context('mytask', args=[
                    Argument('pos', positional=True),
                    Argument('nonpos', default='default')
                ])
                othertask = Context('lolwut')
                result = Parser([mytask]).parse_argv(['mytask', 'lolwut'])
                r = result[0]
                eq_(r.args['pos'].value, 'lolwut')
                eq_(r.args['nonpos'].value, 'default')
                eq_(len(result), 1) # Not 2

            def positional_args_can_still_be_given_as_flags(self):
                # AKA "positional args can come anywhere in the context"
                pos1 = Argument('pos1', positional=True)
                pos2 = Argument('pos2', positional=True)
                nonpos = Argument('nonpos', positional=False, default='lol')
                mytask = Context('mytask', args=[pos1, pos2, nonpos])
                eq_(mytask.positional_args, [pos1, pos2])
                r = Parser([mytask]).parse_argv([
                    'mytask',
                    '--nonpos', 'wut',
                    '--pos2', 'pos2val',
                    'pos1val',
                ])[0]
                eq_(r.args['pos1'].value, 'pos1val')
                eq_(r.args['pos2'].value, 'pos2val')
                eq_(r.args['nonpos'].value, 'wut')

        class equals_signs:
            def _compare(self, argname, invoke, value):
                c = Context('mytask', args=(Argument(argname, kind=str),))
                r = Parser((c,)).parse_argv(['mytask', invoke])
                eq_(r[0].args[argname].value, value)

            def handles_equals_style_long_flags(self):
                self._compare('foo', '--foo=bar', 'bar')

            def handles_equals_style_short_flags(self):
                self._compare('f', '-f=bar', 'bar')

            def does_not_require_escaping_equals_signs_in_value(self):
                self._compare('f', '-f=biz=baz', 'biz=baz')

        def handles_multiple_boolean_flags_per_context(self):
            c = Context('mytask', args=(
                Argument('foo', kind=bool), Argument('bar', kind=bool)
            ))
            r = Parser([c]).parse_argv(['mytask', '--foo', '--bar'])
            a = r[0].args
            eq_(a.foo.value, True)
            eq_(a.bar.value, True)

    class optional_arg_values:
        def setup(self):
            self.context = Context('mytask', args=(
                Argument('foo', optional=True, default='mydefault'),
            ))
            self.parser = Parser([self.context])

        def _expect(self, argstr, expected):
            result = self.parser.parse_argv(['mytask'] + argstr.split())
            eq_(result[0].args.foo.value, expected)

        def no_value_becomes_True(self):
            self._expect('--foo', True)
            self._expect('-f', True)

        def value_given_gets_preserved_normally(self):
            for argstr in (
                '--foo whatever',
                '--foo=whatever',
                '-f whatever',
                '-f=whatever',
            ):
                self._expect(argstr, 'whatever')

        def not_given_at_all_uses_default_value(self):
            self._expect('', 'mydefault')

        def ambiguity_with_unfilled_posargs(self):
            # mytask --foo value-not-posarg
            skip()

        def ambiguity_with_flaglike_value(self):
            # mytask --foo --lolwut
            skip()

        def ambiguity_with_task_name(self):
            # mytask --foo myothertask
            skip()

class ParseResult_(Spec):
    "ParseResult"
    def setup(self):
        self.context = Context('mytask',
            args=(Argument('foo', kind=str), Argument('bar')))
        argv = ['mytask', '--foo', 'foo-val', '--', 'my', 'remainder']
        self.result = Parser((self.context,)).parse_argv(argv)

    def acts_as_a_list_of_parsed_contexts(self):
        eq_(len(self.result), 1)
        eq_(self.result[0].name, 'mytask')

    def exhibits_remainder_attribute(self):
        eq_(self.result.remainder, 'my remainder')

    def to_dict_returns_parsed_contexts_and_args_as_nested_dicts(self):
        eq_(
            self.result.to_dict(),
            {'mytask': {'foo': 'foo-val', 'bar': None}}
        )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    Account Product - Fiscal Classification module for Odoo
#    Copyright (C) 2014-Today GRAP (http://www.grap.coop)
#    @author Sylvain LE GAL (https://twitter.com/legalsylvain)
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.exceptions import ValidationError
from openerp.tests.common import TransactionCase

class Tests(TransactionCase):
    """Tests for 'Account Product - Fiscal Classification' Module"""

    def setUp(self):
        super(Tests, self).setUp()
        self.template_obj = self.env['product.template']
        self.classification_template_obj =\
            self.env['account.product.fiscal.classification.template']
        self.classification_obj =\
            self.env['account.product.fiscal.classification']
        self.wizard_template_obj =\
            self.env['wizard.account.product.fiscal.classification']
        self.wizard_obj = self.env['wizard.change.fiscal.classification']
        self.main_company_id = self.ref('base.main_company')
        self.classification_template_1_id = self.ref(
            'account_product_fiscal_classification.classification_template_1')
        self.classification_1_id = self.ref(
            'account_product_fiscal_classification.classification_1')
        self.classification_2_id = self.ref(
            'account_product_fiscal_classification.classification_2')
        self.template_id = self.ref(
            'account_product_fiscal_classification.product_template_1')
        self.purchase_tax_id = self.ref(
            'account_product_fiscal_classification.account_tax_purchase_1')
        self.sale_tax_1_id = self.ref(
            'account_product_fiscal_classification.account_tax_sale_1')
        self.sale_tax_2_id = self.ref(
            'account_product_fiscal_classification.account_tax_sale_2')

    # Test Section
    def test_01_change_classification(self):
        """Test the behaviour when we change Fiscal Classification for
        products."""
        wizard = self.wizard_obj.create({
            'old_fiscal_classification_id': self.classification_1_id,
            'new_fiscal_classification_id': self.classification_2_id})
        wizard.button_change_fiscal_classification()
        template = self.template_obj.browse(self.template_id)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_2_id,
            "Fiscal Classification change has failed for products via Wizard.")

    def test_02_write_taxes_setting_classification_exist(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination of taxes exist."""
        # Set classification_1 configuration to the product
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, 0, [self.purchase_tax_id]]],
            'taxes_id': [[6, 0, [self.sale_tax_1_id, self.sale_tax_2_id]]],
        }
        template = self.template_obj.create(vals)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_1_id,
            "Recovery of Correct Taxes Group failed during creation.")
        # Set classification_2 configuration to the product
        vals = {
            'supplier_taxes_id': [[6, 0, []]],
            'taxes_id': [[6, 0, [self.sale_tax_2_id]]],
        }
        template.write(vals)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_2_id,
            "Recovery of Correct Taxes Group failed during update.")

    def test_03_write_taxes_setting_classification_doesnt_exist_single(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination doesn't exist.
        (Single Tax)"""
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, 0, [self.purchase_tax_id]]],
            'taxes_id': [[6, 0, [self.sale_tax_1_id]]],
        }
        count_before = self.classification_obj.search_count([])
        self.template_obj.create(vals)
        count_after = self.classification_obj.search_count([])
        self.assertEqual(
            count_before + 1, count_after,
            "New combination must create new Fiscal Classification.")

    def test_04_write_taxes_setting_classification_doesnt_exist_multi(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination doesn't exist.
        (Multiple Taxes)"""
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, False, []]],
            'taxes_id': [[6, False, [self.sale_tax_1_id, self.sale_tax_2_id]]],
        }
        count_before = self.classification_obj.search_count([])
        self.template_obj.create(vals)
        count_after = self.classification_obj.search_count([])
        self.assertEqual(
            count_before + 1, count_after,
            "New combination must create new Fiscal Classification.")

    def test_05_update_fiscal_classification(self):
        """Test if changing a Configuration of a Fiscal Classificationchange
            the product."""
        tg = self.classification_obj.browse([self.classification_1_id])
        tg.write({'customer_tax_ids': [[6, 0, [self.sale_tax_1_id]]]})
        template = self.template_obj.browse([self.template_id])[0]
        self.assertEqual(
            [
                [x.id for x in template.taxes_id],
                [x.id for x in template.supplier_taxes_id]],
            [[self.sale_tax_1_id], [self.purchase_tax_id]],
            "Update taxes in Fiscal Classification must update associated"
            " Products.")

    def test_06_unlink_fiscal_classification(self):
        """Test if unlinking a Fiscal Classification with products fails."""
        classification = self.classification_obj.browse(
            [self.classification_1_id])
        with self.assertRaises(ValidationError):
            classification.unlink()

    def test_07_classification_generate_from_template(self):
        """Test wizard generate fiscal classification from template."""
        wizard_template = self.wizard_template_obj.create({})
        wizard_template.action_create()
        template = self.classification_template_obj.browse(
            self.classification_template_1_id)
        self.assertTrue(self.classification_obj.search(
            [('code', '=', template.code)]))
<commit_msg>fixed wrong xml_id ref -> tests 1,2,3 and 4 pass. 5 and 6 fail.<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    Account Product - Fiscal Classification module for Odoo
#    Copyright (C) 2014-Today GRAP (http://www.grap.coop)
#    @author Sylvain LE GAL (https://twitter.com/legalsylvain)
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 of the
#    License, or (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from openerp.exceptions import ValidationError
from openerp.tests.common import TransactionCase

class Tests(TransactionCase):
    """Tests for 'Account Product - Fiscal Classification' Module"""

    def setUp(self):
        super(Tests, self).setUp()
        self.template_obj = self.env['product.template']
        self.classification_template_obj =\
            self.env['account.product.fiscal.classification.template']
        self.classification_obj =\
            self.env['account.product.fiscal.classification']
        self.wizard_template_obj =\
            self.env['wizard.account.product.fiscal.classification']
        self.wizard_obj = self.env['wizard.change.fiscal.classification']
        self.main_company_id = self.ref('base.main_company')
        self.classification_template_1_id = self.ref(
            'account_product_fiscal_classification.classification_template_1')
        self.classification_1_id = self.ref(
            'account_product_fiscal_classification.fiscal_classification_1')
        self.classification_2_id = self.ref(
            'account_product_fiscal_classification.fiscal_classification_2')
        self.template_id = self.ref(
            'account_product_fiscal_classification.product_template_1')
        self.purchase_tax_id = self.ref(
            'account_product_fiscal_classification.account_tax_purchase_1')
        self.sale_tax_1_id = self.ref(
            'account_product_fiscal_classification.account_tax_sale_1')
        self.sale_tax_2_id = self.ref(
            'account_product_fiscal_classification.account_tax_sale_2')

    # Test Section
    def test_01_change_classification(self):
        """Test the behaviour when we change Fiscal Classification for
        products."""
        wizard = self.wizard_obj.create({
            'old_fiscal_classification_id': self.classification_1_id,
            'new_fiscal_classification_id': self.classification_2_id})
        wizard.button_change_fiscal_classification()
        template = self.template_obj.browse(self.template_id)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_2_id,
            "Fiscal Classification change has failed for products via Wizard.")

    def test_02_write_taxes_setting_classification_exist(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination of taxes exist."""
        # Set classification_1 configuration to the product
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, 0, [self.purchase_tax_id]]],
            'taxes_id': [[6, 0, [self.sale_tax_1_id, self.sale_tax_2_id]]],
        }
        template = self.template_obj.create(vals)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_1_id,
            "Recovery of Correct Taxes Group failed during creation.")
        # Set classification_2 configuration to the product
        vals = {
            'supplier_taxes_id': [[6, 0, []]],
            'taxes_id': [[6, 0, [self.sale_tax_2_id]]],
        }
        template.write(vals)
        self.assertEqual(
            template.fiscal_classification_id.id, self.classification_2_id,
            "Recovery of Correct Taxes Group failed during update.")

    def test_03_write_taxes_setting_classification_doesnt_exist_single(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination doesn't exist.
        (Single Tax)"""
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, 0, [self.purchase_tax_id]]],
            'taxes_id': [[6, 0, [self.sale_tax_1_id]]],
        }
        count_before = self.classification_obj.search_count([])
        self.template_obj.create(vals)
        count_after = self.classification_obj.search_count([])
        self.assertEqual(
            count_before + 1, count_after,
            "New combination must create new Fiscal Classification.")

    def test_04_write_taxes_setting_classification_doesnt_exist_multi(self):
        """Test the behaviour of the function product.template
        write_taxes_setting() when the combination doesn't exist.
        (Multiple Taxes)"""
        vals = {
            'name': 'Product Product Name',
            'company_id': self.main_company_id,
            'supplier_taxes_id': [[6, False, []]],
            'taxes_id': [[6, False, [self.sale_tax_1_id, self.sale_tax_2_id]]],
        }
        count_before = self.classification_obj.search_count([])
        self.template_obj.create(vals)
        count_after = self.classification_obj.search_count([])
        self.assertEqual(
            count_before + 1, count_after,
            "New combination must create new Fiscal Classification.")

    def test_05_update_fiscal_classification(self):
        """Test if changing a Configuration of a Fiscal Classificationchange
            the product."""
        tg = self.classification_obj.browse([self.classification_1_id])
        tg.write({'customer_tax_ids': [[6, 0, [self.sale_tax_1_id]]]})
        template = self.template_obj.browse([self.template_id])[0]
        self.assertEqual(
            [
                [x.id for x in template.taxes_id],
                [x.id for x in template.supplier_taxes_id]],
            [[self.sale_tax_1_id], [self.purchase_tax_id]],
            "Update taxes in Fiscal Classification must update associated"
            " Products.")

    def test_06_unlink_fiscal_classification(self):
        """Test if unlinking a Fiscal Classification with products fails."""
        classification = self.classification_obj.browse(
            [self.classification_1_id])
        with self.assertRaises(ValidationError):
            classification.unlink()

    def test_07_classification_generate_from_template(self):
        """Test wizard generate fiscal classification from template."""
        wizard_template = self.wizard_template_obj.create({})
        wizard_template.action_create()
        template = self.classification_template_obj.browse(
            self.classification_template_1_id)
        self.assertTrue(self.classification_obj.search(
            [('code', '=', template.code)]))
<|endoftext|>