<commit_before>import os, re, fnmatch
from collections import defaultdict
from igf_data.illumina.samplesheet import SampleSheet
import pandas as pd
from igf_data.igfdb.baseadaptor import BaseAdaptor
from igf_data.igfdb.platformadaptor import PlatformAdaptor
from igf_data.igfdb.projectadaptor import ProjectAdaptor
from igf_data.igfdb.seqrunadaptor import SeqrunAdaptor
from igf_data.igfdb.sampleadaptor import SampleAdaptor
from igf_data.igfdb.experimentadaptor import ExperimentAdaptor
from igf_data.igfdb.runadaptor import RunAdaptor
from igf_data.igfdb.collectionadaptor import CollectionAdaptor
from igf_data.igfdb.fileadaptor import FileAdaptor
from igf_data.utils.fileutils import calculate_file_checksum

class Collect_seqrun_fastq_to_db:
  def __init__(self,fastq_dir,model_name,seqrun_igf_id,session_class,\
               samplesheet_filename='SampleSheet.csv',collection_type='demultiplexed_fastq',\
               file_location='HPC_PROJECT'):
    self.fastq_dir=fastq_dir
    self.samplesheet_filename=samplesheet_filename
    self.seqrun_igf_id=seqrun_igf_id
    self.model_name=model_name
    self.session_class=session_class
    self.collection_type=collection_type
    self.file_location=file_location

  def find_fastq_and_build_db_collection(self):
    '''
    A method for finding fastq files and samplesheet under a run directory
    and loading the new files to db with their experiment and run information
    It calculates following entries:
    library_name: Same as sample_id unless mentioned in 'Description' field of
                  samplesheet
    experiment_igf_id: library_name combined with the platform name
                       same library sequenced in different platform will be added
                       as separate experiemnt
    run_igf_id: experiment_igf_id combined with sequencing flowcell_id and lane_id
    collection name: Same as run_igf_id, fastq files will be added to db collection
                     using this id
    collection type: Default type for fastq file collections are 'demultiplexed_fastq'
    file_location: default value is 'HPC_PROJECT'
    '''
    try:
      fastq_files_list=self._collect_fastq_and_sample_info()
      self._build_and_store_exp_run_and_collection_in_db(fastq_files_list=fastq_files_list)
    except:
      raise

  def _get_fastq_and_samplesheet(self):
    fastq_dir=self.fastq_dir
    samplesheet_filename=self.samplesheet_filename
    r1_fastq_regex=re.compile(r'\S+_R1_\d+\.fastq(\.gz)?', re.IGNORECASE)
    r2_fastq_regex=re.compile(r'\S+_R2_\d+\.fastq(\.gz)?', re.IGNORECASE)
    samplesheet_list=list()
    r1_fastq_list=list()
    r2_fastq_list=list()
    for root, dirs, files in os.walk(top=fastq_dir, topdown=True):
      if samplesheet_filename in files:
        samplesheet_list.append(os.path.join(root,samplesheet_filename))
        for file in files:
          if not fnmatch.fnmatch(file, 'Undetermined_'):
            if r1_fastq_regex.match(file):
              r1_fastq_list.append(os.path.join(root,file))
            elif r2_fastq_regex.match(file):
              r2_fastq_list.append(os.path.join(root,file))

    if len(r2_fastq_list) > 0 and len(r1_fastq_list) != len(r2_fastq_list):
      raise ValueError('R1 {0} and R2 {1}'.format(len(r1_fastq_list),\
                                                  len(r2_fastq_list)))

    if len(samplesheet_list) > 1:
      raise ValueError('Found more than one samplesheet file for fastq dir {0}'.\
                       format(fastq_dir))

    return samplesheet_list[0], r1_fastq_list, r2_fastq_list

  def _link_fastq_file_to_sample(self,sample_name,r1_fastq_list, r2_fastq_list):
    sample_files=defaultdict(lambda: defaultdict(lambda: defaultdict()))
    r1_regex=re.compile(sample_name+'_S\d+_L(\d+)_R1_\d+\.fastq(\.gz)?',\
                        re.IGNORECASE)
    for file1 in r1_fastq_list:
      if r1_regex.match(os.path.basename(file1)):
        m=r1_regex.match(os.path.basename(file1))
        lane_id=m.group(1).strip('0')
        sample_files[lane_id]['R1']=file1

    if len(r2_fastq_list) > 0:
      r2_regex=re.compile(sample_name+'_S\d+_L(\d+)_R2_\d+\.fastq(\.gz)?',\
                          re.IGNORECASE)
      for file2 in r2_fastq_list:
        if r2_regex.match(os.path.basename(file2)):
          m=r2_regex.match(os.path.basename(file2))
          lane_id=m.group(1).strip('0')
          sample_files[lane_id]['R2']=file2
  return sample_files

  def _collect_fastq_and_sample_info(self):
    fastq_dir=self.fastq_dir
    samplesheet_filename=self.samplesheet_filename
    seqrun_igf_id=self.seqrun_igf_id
    model_name=self.model_name
    (samplesheet_file, r1_fastq_list, r2_fastq_list)=\
        self._get_fastq_and_samplesheet()
    samplesheet_data=SampleSheet(infile=samplesheet_file)
    fastq_files_list=list()
    for row in samplesheet_data._data:
      sample_name=row['Sample_Name']
      sample_id=row['Sample_ID']
      project_name=row['Sample_Project']
      description=row['Description']
      sample_files=self._link_fastq_file_to_sample(sample_name,r1_fastq_list, \
                                                   r2_fastq_list)
      for lane, lane_files in sample_files.items():
        fastq_info={'sample_igf_id':sample_id,
                    'sample_name':sample_name,
                    'project_igf_id':project_name,
                    'lane_number':lane,
                    'seqrun_igf_id':seqrun_igf_id,
                    'platform_name':model_name,
                    'flowcell_id':flowcell_id,
                    'description':description
                   }
        for read_type, filepath in lane_files.items():
          fastq_info.update({read_type:filepath})                               # allowing only one file per lane per read type
          fastq_files_list.append(fastq_info)                                   # adding entries per sample per lane
    return fastq_files_list

  def _calculate_experiment_run_and_file_info(self,data,restricted_list):
    if not isinstance(data, pd.Series):
      data=pd.Series(data)
    # set library id
    library_id=None
    if data.description and data.description not in restricted_list:
      library_id=data.description
    else:
      library_id=data.sample_igf_id

    # calcaulate experiment id
    experiment_id='{0}_{1}'.format(library_id,data.platform_name)
    data['library_name']=library_id
    data['experiment_igf_id']=experiment_id
    # calculate run id
    run_igf_id='{0}_{1}_{2}'.format(experiment_id, \
                                    data.flowcell_id, \
                                    data.lane_number)
    data['run_igf_id']=run_igf_id
    # set collection name and type
    data['name']=run_igf_id
    data['type']==self.collection_type
    data['location']=self.file_location
    # set file md5 and size
    if 'R1' in data:
      data['R1_md5']=calculate_file_checksum(filepath=data.R1, \
                                             hasher='md5')
      data['R1_size']=os.path.getsize(data.R1)
    if 'R2' in data:
      data['R2_md5']=calculate_file_checksum(filepath=data.R2, \
                                             hasher='md5')
      data['R2_size']=os.path.getsize(data.R2)
    # set library strategy
    library_layout='SINGLE'
    if 'R1' in data and 'R2' in data and \
      data.R1 is not None and data.R2 is not None:
      library_layout='PAIRED'

    data['library_layout']=library_layout
    return data

  def _reformat_file_group_data(self,data):
    if isinstance(data, pd.DataFrame):
      data=data.to_dict(orient='records')

    if not isinstance(data,list):
      raise ValueError('Expecting list got {0}'.format(type(data)))

    reformatted_file_group_data=list()
    reformatted_file_data=list()

    for row in data:
      collection_name=None
      collection_type=None
      file_location=None
      if 'name' in row.keys():
        collection_name=row['name']
      if 'type' in row.keys():
        collection_type=row['type']
      if 'location' in row.keys():
        file_location=row['location']
      if 'R1' in row.keys():
        r1_file_path=row['R1']
        r1_file_size=row['R1_size'] if 'R1_size' in row.keys() else None
        r1_file_md5=row['R1_md5'] if 'R1_md5' in row.keys() else None
        reformatted_file_data.append({'file_path':r1_file_path,
                                      'md5':r1_file_md5,
                                      'location':file_location,
                                      'size':r1_file_size})
        reformatted_file_group_data.append({'name':collection_name,
                                            'type':collection_type,
                                            'file_path':r1_file_path})
      if 'R2' in row.keys():
        r2_file_path=row['R2']
        r2_file_size=row['R2_size'] if 'R2_size' in row.keys() else None
        r2_file_md5=row['R2_md5'] if 'R2_md5' in row.keys() else None
        reformatted_file_data.append({'file_path':r2_file_path,
                                      'md5':r2_file_md5,
                                      'location':file_location,
                                      'size':r2_file_size})
        reformatted_file_group_data.append({'name':collection_name,
                                            'type':collection_type,
                                            'file_path':r2_file_path})
    return pd.DataFrame(reformatted_file_data), pd.DataFrame(reformatted_file_group_data)

  def _build_and_store_exp_run_and_collection_in_db(self,fastq_files_list, restricted_list=['10X']):
    session_class=self.session_class
    db_connected=False
    try:
      dataframe=pd.DataFrame(fastq_files_list)
      # calculate additional detail
      dataframe=dataframe.apply(lambda data: \
                                self._calculate_experiment_run_and_file_info(data,
                                                               restricted_list),\
                                axis=1)
      # get file data
      file_group_columns=['name','type','location',
                          'R1','R1_md5','R1_size','R2',
                          'R2_md5','R2_size']
      file_group_data=dataframe.loc[:,file_group_columns]
      file_group_data=file_group_data.drop_duplicates()
      (file_data,file_group_data)=self._reformat_file_group_data(data=file_group_data)
      # get base session
      base=BaseAdaptor(**{'session_class':session_class})
      base.start_session()
      db_connected=True
      # get experiment data
      experiment_columns=base.get_table_columns(table_name=Experiment, \
                                                excluded_columns=['experiment_id',
                                                                  'project_id',
                                                                  'sample_id' ])
      experiment_columns.extend(['project_igf_id',
                                 'sample_igf_id'])
      exp_data=dataframe.loc[:,experiment_columns]
      exp_data=exp_data.drop_duplicates()
      # get run data
      run_columns=base.get_table_columns(table_name=Run, \
                                         excluded_columns=['run_id',
                                                           'seqrun_id',
                                                           'experiment_id',
                                                           'date_created',
                                                           'status'
                                                          ])
      run_columns.extend(['seqrun_igf_id',
                          'experiment_igf_id'])
      run_data=dataframe.loc[:,run_columns]
      run_data=run_data.drop_duplicates()
      # get collection data
      collection_columns=['name',
                          'type']
      collection_data=dataframe.loc[:,collection_columns]
      collection_data=collection_data.drop_duplicates()
      # store experiment to db
      ea=ExperimentAdaptor(**{'session':base.session})
      ea.store_project_and_attribute_data(data=exp_data,autosave=False)
      base.session.flush()
      # store run to db
      ra=RunAdaptor(**{'session':base.session})
      ra.store_run_and_attribute_data(data=run_data,autosave=False)
      base.session.flush()
      # store file to db
      fa=FileAdaptor(**{'session':base.session})
      fa.store_file_and_attribute_data(data=file_data,autosave=False)
      base.session.flush()
      # store collection to db
      ca=CollectionAdaptor(**{'session':base.session})
      ca.store_collection_and_attribute_data(data=collection_data,\
                                             autosave=False)
      base.session.flush()
      ca.create_collection_group(data=file_group_data,autosave=False)
      base.commit_session()
    except:
      if db_connected:
        base.rollback_session()
      raise
    finally:
      if db_connected:
        base.close_session()
      <commit_msg>added return back to function<commit_after>import os, re, fnmatch
from collections import defaultdict
from igf_data.illumina.samplesheet import SampleSheet
import pandas as pd
from igf_data.igfdb.baseadaptor import BaseAdaptor
from igf_data.igfdb.platformadaptor import PlatformAdaptor
from igf_data.igfdb.projectadaptor import ProjectAdaptor
from igf_data.igfdb.seqrunadaptor import SeqrunAdaptor
from igf_data.igfdb.sampleadaptor import SampleAdaptor
from igf_data.igfdb.experimentadaptor import ExperimentAdaptor
from igf_data.igfdb.runadaptor import RunAdaptor
from igf_data.igfdb.collectionadaptor import CollectionAdaptor
from igf_data.igfdb.fileadaptor import FileAdaptor
from igf_data.utils.fileutils import calculate_file_checksum

class Collect_seqrun_fastq_to_db:
  def __init__(self,fastq_dir,model_name,seqrun_igf_id,session_class,\
               samplesheet_filename='SampleSheet.csv',collection_type='demultiplexed_fastq',\
               file_location='HPC_PROJECT'):
    self.fastq_dir=fastq_dir
    self.samplesheet_filename=samplesheet_filename
    self.seqrun_igf_id=seqrun_igf_id
    self.model_name=model_name
    self.session_class=session_class
    self.collection_type=collection_type
    self.file_location=file_location

  def find_fastq_and_build_db_collection(self):
    '''
    A method for finding fastq files and samplesheet under a run directory
    and loading the new files to db with their experiment and run information
    It calculates following entries:
    library_name: Same as sample_id unless mentioned in 'Description' field of
                  samplesheet
    experiment_igf_id: library_name combined with the platform name
                       same library sequenced in different platform will be added
                       as separate experiemnt
    run_igf_id: experiment_igf_id combined with sequencing flowcell_id and lane_id
    collection name: Same as run_igf_id, fastq files will be added to db collection
                     using this id
    collection type: Default type for fastq file collections are 'demultiplexed_fastq'
    file_location: default value is 'HPC_PROJECT'
    '''
    try:
      fastq_files_list=self._collect_fastq_and_sample_info()
      self._build_and_store_exp_run_and_collection_in_db(fastq_files_list=fastq_files_list)
    except:
      raise

  def _get_fastq_and_samplesheet(self):
    fastq_dir=self.fastq_dir
    samplesheet_filename=self.samplesheet_filename
    r1_fastq_regex=re.compile(r'\S+_R1_\d+\.fastq(\.gz)?', re.IGNORECASE)
    r2_fastq_regex=re.compile(r'\S+_R2_\d+\.fastq(\.gz)?', re.IGNORECASE)
    samplesheet_list=list()
    r1_fastq_list=list()
    r2_fastq_list=list()
    for root, dirs, files in os.walk(top=fastq_dir, topdown=True):
      if samplesheet_filename in files:
        samplesheet_list.append(os.path.join(root,samplesheet_filename))
        for file in files:
          if not fnmatch.fnmatch(file, 'Undetermined_'):
            if r1_fastq_regex.match(file):
              r1_fastq_list.append(os.path.join(root,file))
            elif r2_fastq_regex.match(file):
              r2_fastq_list.append(os.path.join(root,file))

    if len(r2_fastq_list) > 0 and len(r1_fastq_list) != len(r2_fastq_list):
      raise ValueError('R1 {0} and R2 {1}'.format(len(r1_fastq_list),\
                                                  len(r2_fastq_list)))

    if len(samplesheet_list) > 1:
      raise ValueError('Found more than one samplesheet file for fastq dir {0}'.\
                       format(fastq_dir))

    return samplesheet_list[0], r1_fastq_list, r2_fastq_list

  def _link_fastq_file_to_sample(self,sample_name,r1_fastq_list, r2_fastq_list):
    sample_files=defaultdict(lambda: defaultdict(lambda: defaultdict()))
    r1_regex=re.compile(sample_name+'_S\d+_L(\d+)_R1_\d+\.fastq(\.gz)?',\
                        re.IGNORECASE)
    for file1 in r1_fastq_list:
      if r1_regex.match(os.path.basename(file1)):
        m=r1_regex.match(os.path.basename(file1))
        lane_id=m.group(1).strip('0')
        sample_files[lane_id]['R1']=file1

    if len(r2_fastq_list) > 0:
      r2_regex=re.compile(sample_name+'_S\d+_L(\d+)_R2_\d+\.fastq(\.gz)?',\
                          re.IGNORECASE)
      for file2 in r2_fastq_list:
        if r2_regex.match(os.path.basename(file2)):
          m=r2_regex.match(os.path.basename(file2))
          lane_id=m.group(1).strip('0')
          sample_files[lane_id]['R2']=file2
    return sample_files

  def _collect_fastq_and_sample_info(self):
    fastq_dir=self.fastq_dir
    samplesheet_filename=self.samplesheet_filename
    seqrun_igf_id=self.seqrun_igf_id
    model_name=self.model_name
    (samplesheet_file, r1_fastq_list, r2_fastq_list)=\
        self._get_fastq_and_samplesheet()
    samplesheet_data=SampleSheet(infile=samplesheet_file)
    fastq_files_list=list()
    for row in samplesheet_data._data:
      sample_name=row['Sample_Name']
      sample_id=row['Sample_ID']
      project_name=row['Sample_Project']
      description=row['Description']
      sample_files=self._link_fastq_file_to_sample(sample_name,r1_fastq_list, \
                                                   r2_fastq_list)
      for lane, lane_files in sample_files.items():
        fastq_info={'sample_igf_id':sample_id,
                    'sample_name':sample_name,
                    'project_igf_id':project_name,
                    'lane_number':lane,
                    'seqrun_igf_id':seqrun_igf_id,
                    'platform_name':model_name,
                    'flowcell_id':flowcell_id,
                    'description':description
                   }
        for read_type, filepath in lane_files.items():
          fastq_info.update({read_type:filepath})                               # allowing only one file per lane per read type
          fastq_files_list.append(fastq_info)                                   # adding entries per sample per lane
    return fastq_files_list

  def _calculate_experiment_run_and_file_info(self,data,restricted_list):
    if not isinstance(data, pd.Series):
      data=pd.Series(data)
    # set library id
    library_id=None
    if data.description and data.description not in restricted_list:
      library_id=data.description
    else:
      library_id=data.sample_igf_id

    # calcaulate experiment id
    experiment_id='{0}_{1}'.format(library_id,data.platform_name)
    data['library_name']=library_id
    data['experiment_igf_id']=experiment_id
    # calculate run id
    run_igf_id='{0}_{1}_{2}'.format(experiment_id, \
                                    data.flowcell_id, \
                                    data.lane_number)
    data['run_igf_id']=run_igf_id
    # set collection name and type
    data['name']=run_igf_id
    data['type']==self.collection_type
    data['location']=self.file_location
    # set file md5 and size
    if 'R1' in data:
      data['R1_md5']=calculate_file_checksum(filepath=data.R1, \
                                             hasher='md5')
      data['R1_size']=os.path.getsize(data.R1)
    if 'R2' in data:
      data['R2_md5']=calculate_file_checksum(filepath=data.R2, \
                                             hasher='md5')
      data['R2_size']=os.path.getsize(data.R2)
    # set library strategy
    library_layout='SINGLE'
    if 'R1' in data and 'R2' in data and \
      data.R1 is not None and data.R2 is not None:
      library_layout='PAIRED'

    data['library_layout']=library_layout
    return data

  def _reformat_file_group_data(self,data):
    if isinstance(data, pd.DataFrame):
      data=data.to_dict(orient='records')

    if not isinstance(data,list):
      raise ValueError('Expecting list got {0}'.format(type(data)))

    reformatted_file_group_data=list()
    reformatted_file_data=list()

    for row in data:
      collection_name=None
      collection_type=None
      file_location=None
      if 'name' in row.keys():
        collection_name=row['name']
      if 'type' in row.keys():
        collection_type=row['type']
      if 'location' in row.keys():
        file_location=row['location']
      if 'R1' in row.keys():
        r1_file_path=row['R1']
        r1_file_size=row['R1_size'] if 'R1_size' in row.keys() else None
        r1_file_md5=row['R1_md5'] if 'R1_md5' in row.keys() else None
        reformatted_file_data.append({'file_path':r1_file_path,
                                      'md5':r1_file_md5,
                                      'location':file_location,
                                      'size':r1_file_size})
        reformatted_file_group_data.append({'name':collection_name,
                                            'type':collection_type,
                                            'file_path':r1_file_path})
      if 'R2' in row.keys():
        r2_file_path=row['R2']
        r2_file_size=row['R2_size'] if 'R2_size' in row.keys() else None
        r2_file_md5=row['R2_md5'] if 'R2_md5' in row.keys() else None
        reformatted_file_data.append({'file_path':r2_file_path,
                                      'md5':r2_file_md5,
                                      'location':file_location,
                                      'size':r2_file_size})
        reformatted_file_group_data.append({'name':collection_name,
                                            'type':collection_type,
                                            'file_path':r2_file_path})
    return pd.DataFrame(reformatted_file_data), pd.DataFrame(reformatted_file_group_data)

  def _build_and_store_exp_run_and_collection_in_db(self,fastq_files_list, restricted_list=['10X']):
    session_class=self.session_class
    db_connected=False
    try:
      dataframe=pd.DataFrame(fastq_files_list)
      # calculate additional detail
      dataframe=dataframe.apply(lambda data: \
                                self._calculate_experiment_run_and_file_info(data,
                                                               restricted_list),\
                                axis=1)
      # get file data
      file_group_columns=['name','type','location',
                          'R1','R1_md5','R1_size','R2',
                          'R2_md5','R2_size']
      file_group_data=dataframe.loc[:,file_group_columns]
      file_group_data=file_group_data.drop_duplicates()
      (file_data,file_group_data)=self._reformat_file_group_data(data=file_group_data)
      # get base session
      base=BaseAdaptor(**{'session_class':session_class})
      base.start_session()
      db_connected=True
      # get experiment data
      experiment_columns=base.get_table_columns(table_name=Experiment, \
                                                excluded_columns=['experiment_id',
                                                                  'project_id',
                                                                  'sample_id' ])
      experiment_columns.extend(['project_igf_id',
                                 'sample_igf_id'])
      exp_data=dataframe.loc[:,experiment_columns]
      exp_data=exp_data.drop_duplicates()
      # get run data
      run_columns=base.get_table_columns(table_name=Run, \
                                         excluded_columns=['run_id',
                                                           'seqrun_id',
                                                           'experiment_id',
                                                           'date_created',
                                                           'status'
                                                          ])
      run_columns.extend(['seqrun_igf_id',
                          'experiment_igf_id'])
      run_data=dataframe.loc[:,run_columns]
      run_data=run_data.drop_duplicates()
      # get collection data
      collection_columns=['name',
                          'type']
      collection_data=dataframe.loc[:,collection_columns]
      collection_data=collection_data.drop_duplicates()
      # store experiment to db
      ea=ExperimentAdaptor(**{'session':base.session})
      ea.store_project_and_attribute_data(data=exp_data,autosave=False)
      base.session.flush()
      # store run to db
      ra=RunAdaptor(**{'session':base.session})
      ra.store_run_and_attribute_data(data=run_data,autosave=False)
      base.session.flush()
      # store file to db
      fa=FileAdaptor(**{'session':base.session})
      fa.store_file_and_attribute_data(data=file_data,autosave=False)
      base.session.flush()
      # store collection to db
      ca=CollectionAdaptor(**{'session':base.session})
      ca.store_collection_and_attribute_data(data=collection_data,\
                                             autosave=False)
      base.session.flush()
      ca.create_collection_group(data=file_group_data,autosave=False)
      base.commit_session()
    except:
      if db_connected:
        base.rollback_session()
      raise
    finally:
      if db_connected:
        base.close_session()
      <|endoftext|>


<commit_before>
# -*- coding: utf-8 -*-

from model.group import Contact

def test_modify_contact_firstname(app):
    if app.contact.count() == 0:
        app.contact.fill_new(Contact(lastname="tester"))
    old_contact_list = app.contact.get_contact_list()
    contact = Contact(firstname="tretret")
    contact.id = old_contact_list[0].id
    app.contact.modify_first_contact(contact)
    assert len(old_contact_list) ==  app.contact.count()
    new_contact = app.contact.get_contact_list()
    old_contact_list[0] = contact
    assert sorted(old_contact_list, key=Contact.id_or_max) == sorted(new_contact, key=Contact.id_or_max)

#def test_test_modify_contact_address(app):
#    if app.contact.count() ==0:
#        app.contact.fill_new(Contact(company="tester"))
#   old_contact = app.contact.get_contact_list()
#    app.contact.modify_first_contact(Contact(address="Moscow"))
#    new_contact = app.contact.get_contact_list()
 #   assert len(old_contact) == len(new_contact)
<commit_msg>тест модификации конт+кэш+проверка по длине<commit_after>
# -*- coding: utf-8 -*-

from model.group import Contact

def test_modify_contact_firstname(app):
    if app.contact.count() == 0:
        app.contact.fill_new(Contact(lastname="tester"))
    old_contact_list = app.contact.get_contact_list()
    contact = Contact(firstname="tretret")
    contact.id = old_contact_list[0].id
    app.contact.modify_first_contact(contact)
    assert len(old_contact_list) ==  app.contact.count()
    new_contact_list = app.contact.get_contact_list()
    old_contact_list[0] = contact
    assert sorted(old_contact_list, key=Contact.id_or_max) == sorted(new_contact_list, key=Contact.id_or_max)

#def test_test_modify_contact_address(app):
#    if app.contact.count() ==0:
#        app.contact.fill_new(Contact(company="tester"))
#   old_contact = app.contact.get_contact_list()
#    app.contact.modify_first_contact(Contact(address="Moscow"))
#    new_contact = app.contact.get_contact_list()
 #   assert len(old_contact) == len(new_contact)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>). All Rights Reserved
#    $Id$
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU 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 General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from tools import flatten, reverse_enumerate
import fields

class expression(object):
    """
    parse a domain expression
    use a real polish notation
    leafs are still in a ('foo', '=', 'bar') format
    For more info: http://christophe-simonis-at-tiny.blogspot.com/2008/08/new-new-domain-notation.html
    """

    def _is_operator(self, element):
        return isinstance(element, (str, unicode)) and element in ['&', '|', '!']

    def _is_leaf(self, element, internal=False):
        OPS = ('=', '!=', '<>', '<=', '<', '>', '>=', '=like', 'like', 'not like', 'ilike', 'not ilike', 'in', 'not in', 'child_of')
        INTERNAL_OPS = OPS + ('inselect',)
        return (isinstance(element, tuple) or isinstance(element, list)) \
           and len(element) == 3 \
           and (((not internal) and element[1] in OPS) \
                or (internal and element[1] in INTERNAL_OPS))

    def __execute_recursive_in(self, cr, s, f, w, ids):
        res = []
        for i in range(0, len(ids), cr.IN_MAX):
            subids = ids[i:i+cr.IN_MAX]
            cr.execute('SELECT "%s"'    \
                       '  FROM "%s"'    \
                       ' WHERE "%s" in (%s)' % (s, f, w, ','.join(['%s']*len(subids))),
                       subids)
            res.extend([r[0] for r in cr.fetchall()])
        return res

    def __init__(self, exp):
        # check if the expression is valid
        if not reduce(lambda acc, val: acc and (self._is_operator(val) or self._is_leaf(val)), exp, True):
            raise ValueError('Bad domain expression: %r' % (exp,))
        self.__exp = exp
        self.__tables = {}  # used to store the table to use for the sql generation. key = index of the leaf
        self.__joins = []
        self.__main_table = None # 'root' table. set by parse()
        self.__DUMMY_LEAF = (1, '=', 1) # a dummy leaf that must not be parsed or sql generated

    def parse(self, cr, uid, table, context):
        """ transform the leafs of the expression """
        if not self.__exp:
            return self

        def _rec_get(ids, table, parent=None, left='id', prefix=''):
            if table._parent_store and (not table.pool._init):
# TODO: Improve where joins are implemented for many with '.', replace by:
# doms += ['&',(prefix+'.parent_left','<',o.parent_right),(prefix+'.parent_left','>=',o.parent_left)]
                doms = []
                for o in table.browse(cr, uid, ids, context=context):
                    if doms:
                        doms.insert(0, '|')
                    doms += ['&', ('parent_left', '<', o.parent_right), ('parent_left', '>=', o.parent_left)]
                if prefix:
                    return [(left, 'in', table.search(cr, uid, doms, context=context))]
                return doms
            else:
                def rg(ids, table, parent):
                    if not ids:
                        return []
                    ids2 = table.search(cr, uid, [(parent, 'in', ids)], context=context)
                    return ids + rg(ids2, table, parent)
                return [(left, 'in', rg(ids, table, parent or table._parent_name))]

        self.__main_table = table

        i = -1
        while i + 1<len(self.__exp):
            i += 1
            e = self.__exp[i]
            if self._is_operator(e) or e == self.__DUMMY_LEAF:
                continue
            left, operator, right = e

            working_table = table
            main_table = table
            fargs = left.split('.', 1)
            index = i
            if left in table._inherit_fields:
                while True:
                    field = main_table._columns.get(fargs[0], False)
                    if field:
                        working_table = main_table
                        self.__tables[i] = working_table
                        break
                    working_table = main_table.pool.get(main_table._inherit_fields[left][0])
                    if working_table not in self.__tables.values():
                        self.__joins.append(('%s.%s=%s.%s' % (working_table._table, 'id', main_table._table, main_table._inherits[working_table._name]), working_table._table))
                        self.__tables[index] = working_table
                        index += 1
                    main_table = working_table

            field = working_table._columns.get(fargs[0], False)
            if not field:
                if left == 'id' and operator == 'child_of':
                    dom = _rec_get(right, working_table)
                    self.__exp = self.__exp[:i] + dom + self.__exp[i+1:]
                continue

            field_obj = table.pool.get(field._obj)
            if len(fargs) > 1:
                if field._type == 'many2one':
                    right = field_obj.search(cr, uid, [(fargs[1], operator, right)], context=context)
                    self.__exp[i] = (fargs[0], 'in', right)
                continue

            if field._properties:
                # this is a function field
                if not field.store:
                    if not field._fnct_search:
                        # the function field doesn't provide a search function and doesn't store
                        # values in the database, so we must ignore it : we generate a dummy leaf
                        self.__exp[i] = self.__DUMMY_LEAF
                    else:
                        subexp = field.search(cr, uid, table, left, [self.__exp[i]])
                        # we assume that the expression is valid
                        # we create a dummy leaf for forcing the parsing of the resulting expression
                        self.__exp[i] = '&'
                        self.__exp.insert(i + 1, self.__DUMMY_LEAF)
                        for j, se in enumerate(subexp):
                            self.__exp.insert(i + 2 + j, se)

                # else, the value of the field is store in the database, so we search on it

            elif field._type == 'one2many':
                if isinstance(right, basestring):
                    ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], operator, limit=None)]
                else:
                    ids2 = list(right)
                if not ids2:
                    self.__exp[i] = ('id', '=', '0')
                else:
                    self.__exp[i] = ('id', 'in', self.__execute_recursive_in(cr, field._fields_id, field_obj._table, 'id', ids2))

            elif field._type == 'many2many':
                #FIXME
                if operator == 'child_of':
                    if isinstance(right, basestring):
                        ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], 'like', limit=None)]
                    else:
                        ids2 = list(right)

                    def _rec_convert(ids):
                        if field_obj == table:
                            return ids
                        return self.__execute_recursive_in(cr, field._id1, field._rel, field._id2, ids)

                    dom = _rec_get(ids2, field_obj)
                    ids2 = field_obj.search(cr, uid, dom, context=context)
                    self.__exp[i] = ('id', 'in', _rec_convert(ids2))
                else:
                    if isinstance(right, basestring):
                        res_ids = [x[0] for x in field_obj.name_search(cr, uid, right, [], operator)]
                    else:
                        res_ids = list(right)
                    self.__exp[i] = ('id', 'in', self.__execute_recursive_in(cr, field._id1, field._rel, field._id2, res_ids) or [0])
            elif field._type == 'many2one':
                if operator == 'child_of':
                    if isinstance(right, basestring):
                        ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], 'like', limit=None)]
                    else:
                        ids2 = list(right)

                    self.__operator = 'in'
                    if field._obj != working_table._name:
                        dom = _rec_get(ids2, field_obj, left=left, prefix=field._obj)
                    else:
                        dom = _rec_get(ids2, working_table, parent=left)
                    self.__exp = self.__exp[:i] + dom + self.__exp[i+1:]
                else:
                    if isinstance(right, basestring): # and not isinstance(field, fields.related):
                        c = context.copy()
                        c['active_test'] = False
                        res_ids = field_obj.name_search(cr, uid, right, [], operator, limit=None, context=c)
                        right = map(lambda x: x[0], res_ids)
                        self.__exp[i] = (left, 'in', right)
            else:
                # other field type
                # add the time part to datetime field when it's not there:
                if field._type == 'datetime' and self.__exp[i][2] and len(self.__exp[i][2]) == 10:

                    self.__exp[i] = list(self.__exp[i])

                    if operator in ('>', '>='):
                        self.__exp[i][2] += ' 00:00:00'
                    elif operator in ('<', '<='):
                        self.__exp[i][2] += ' 23:59:59'

                    self.__exp[i] = tuple(self.__exp[i])

                if field.translate:
                    if operator in ('like', 'ilike', 'not like', 'not ilike'):
                        right = '%%%s%%' % right

                    operator = operator == '=like' and 'like' or operator

                    query1 = '( SELECT res_id'          \
                             '    FROM ir_translation'  \
                             '   WHERE name = %s'       \
                             '     AND lang = %s'       \
                             '     AND type = %s'
                    instr = ' %s'
                    #Covering in,not in operators with operands (%s,%s) ,etc.
                    if operator in ['in','not in']:
                        instr = ','.join(['%s'] * len(right))
                        query1 += '     AND value ' + operator +  ' ' +" (" + instr + ")"   \
                             ') UNION ('                \
                             '  SELECT id'              \
                             '    FROM "' + working_table._table + '"'       \
                             '   WHERE "' + left + '" ' + operator + ' ' +" (" + instr + "))"
                    else:
                        query1 += '     AND value ' + operator + instr +   \
                             ') UNION ('                \
                             '  SELECT id'              \
                             '    FROM "' + working_table._table + '"'       \
                             '   WHERE "' + left + '" ' + operator + instr + ")"

                    query2 = [working_table._name + ',' + left,
                              context.get('lang', False) or 'en_US',
                              'model',
                              right,
                              right,
                             ]

                    self.__exp[i] = ('id', 'inselect', (query1, query2))

        return self

    def __leaf_to_sql(self, leaf, table):
        if leaf == self.__DUMMY_LEAF:
            return ('(1=1)', [])
        left, operator, right = leaf

        if operator == 'inselect':
            query = '(%s.%s in (%s))' % (table._table, left, right[0])
            params = right[1]
        elif operator in ['in', 'not in']:
            params = right[:]
            len_before = len(params)
            for i in range(len_before)[::-1]:
                if params[i] == False:
                    del params[i]

            len_after = len(params)
            check_nulls = len_after != len_before
            query = '(1=0)'

            if len_after:
                if left == 'id':
                    instr = ','.join(['%s'] * len_after)
                else:
                    instr = ','.join([table._columns[left]._symbol_set[0]] * len_after)
                query = '(%s.%s %s (%s))' % (table._table, left, operator, instr)

            if check_nulls:
                query = '(%s OR %s.%s IS NULL)' % (query, table._table, left)
        else:
            params = []

            if right == False and (leaf[0] in table._columns)  and table._columns[leaf[0]]._type=="boolean"  and (operator == '='):
                query = '(%s.%s IS NULL or %s.%s = false )' % (table._table, left,table._table, left)
            elif (((right == False) and (type(right)==bool)) or (right is None)) and (operator == '='):
                query = '%s.%s IS NULL ' % (table._table, left)
            elif right == False and (leaf[0] in table._columns)  and table._columns[leaf[0]]._type=="boolean"  and (operator in ['<>', '!=']):
                query = '(%s.%s IS NOT NULL and %s.%s != false)' % (table._table, left,table._table, left)
            elif (((right == False) and (type(right)==bool)) or right is None) and (operator in ['<>', '!=']):
                query = '%s.%s IS NOT NULL' % (table._table, left)
            else:
                if left == 'id':
                    query = '%s.id %s %%s' % (table._table, operator)
                    params = right
                else:
                    like = operator in ('like', 'ilike', 'not like', 'not ilike')

                    op = operator == '=like' and 'like' or operator
                    if left in table._columns:
                        format = like and '%s' or table._columns[left]._symbol_set[0]
                        query = '(%s.%s %s %s)' % (table._table, left, op, format)
                    else:
                        query = "(%s.%s %s '%s')" % (table._table, left, op, right)

                    add_null = False
                    if like:
                        if isinstance(right, str):
                            str_utf8 = right
                        elif isinstance(right, unicode):
                            str_utf8 = right.encode('utf-8')
                        else:
                            str_utf8 = str(right)
                        params = '%%%s%%' % str_utf8
                        add_null = not str_utf8
                    elif left in table._columns:
                        params = table._columns[left]._symbol_set[1](right)

                    if add_null:
                        query = '(%s OR %s IS NULL)' % (query, left)

        if isinstance(params, basestring):
            params = [params]
        return (query, params)

    def to_sql(self):
        stack = []
        params = []
        for i, e in reverse_enumerate(self.__exp):
            if self._is_leaf(e, internal=True):
                table = self.__tables.get(i, self.__main_table)
                q, p = self.__leaf_to_sql(e, table)
                params.insert(0, p)
                stack.append(q)
            else:
                if e == '!':
                    stack.append('(NOT (%s))' % (stack.pop(),))
                else:
                    ops = {'&': ' AND ', '|': ' OR '}
                    q1 = stack.pop()
                    q2 = stack.pop()
                    stack.append('(%s %s %s)' % (q1, ops[e], q2,))

        query = ' AND '.join(reversed(stack))
        joins = ' AND '.join(map(lambda j: j[0], self.__joins))
        if joins:
            query = '(%s) AND (%s)' % (joins, query)
        return (query, flatten(params))

    def get_tables(self):
        return ['"%s"' % t._table for t in set(self.__tables.values()+[self.__main_table])]

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:

<commit_msg>[FIX] recursive child_of on one2many and many2many<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-
##############################################################################
#
#    OpenERP, Open Source Management Solution
#    Copyright (C) 2004-2009 Tiny SPRL (<http://tiny.be>). All Rights Reserved
#    $Id$
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU 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 General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
##############################################################################

from tools import flatten, reverse_enumerate
import fields

class expression(object):
    """
    parse a domain expression
    use a real polish notation
    leafs are still in a ('foo', '=', 'bar') format
    For more info: http://christophe-simonis-at-tiny.blogspot.com/2008/08/new-new-domain-notation.html
    """

    def _is_operator(self, element):
        return isinstance(element, (str, unicode)) and element in ['&', '|', '!']

    def _is_leaf(self, element, internal=False):
        OPS = ('=', '!=', '<>', '<=', '<', '>', '>=', '=like', 'like', 'not like', 'ilike', 'not ilike', 'in', 'not in', 'child_of')
        INTERNAL_OPS = OPS + ('inselect',)
        return (isinstance(element, tuple) or isinstance(element, list)) \
           and len(element) == 3 \
           and (((not internal) and element[1] in OPS) \
                or (internal and element[1] in INTERNAL_OPS))

    def __execute_recursive_in(self, cr, s, f, w, ids):
        res = []
        for i in range(0, len(ids), cr.IN_MAX):
            subids = ids[i:i+cr.IN_MAX]
            cr.execute('SELECT "%s"'    \
                       '  FROM "%s"'    \
                       ' WHERE "%s" in (%s)' % (s, f, w, ','.join(['%s']*len(subids))),
                       subids)
            res += [r[0] for r in cr.fetchall() if r[0]]
        return res + (res and self.__execute_recursive_in(cr, s, f, w, res) or [])

    def __init__(self, exp):
        # check if the expression is valid
        if not reduce(lambda acc, val: acc and (self._is_operator(val) or self._is_leaf(val)), exp, True):
            raise ValueError('Bad domain expression: %r' % (exp,))
        self.__exp = exp
        self.__tables = {}  # used to store the table to use for the sql generation. key = index of the leaf
        self.__joins = []
        self.__main_table = None # 'root' table. set by parse()
        self.__DUMMY_LEAF = (1, '=', 1) # a dummy leaf that must not be parsed or sql generated

    def parse(self, cr, uid, table, context):
        """ transform the leafs of the expression """
        if not self.__exp:
            return self

        def _rec_get(ids, table, parent=None, left='id', prefix=''):
            if table._parent_store and (not table.pool._init):
# TODO: Improve where joins are implemented for many with '.', replace by:
# doms += ['&',(prefix+'.parent_left','<',o.parent_right),(prefix+'.parent_left','>=',o.parent_left)]
                doms = []
                for o in table.browse(cr, uid, ids, context=context):
                    if doms:
                        doms.insert(0, '|')
                    doms += ['&', ('parent_left', '<', o.parent_right), ('parent_left', '>=', o.parent_left)]
                if prefix:
                    return [(left, 'in', table.search(cr, uid, doms, context=context))]
                return doms
            else:
                def rg(ids, table, parent):
                    if not ids:
                        return []
                    ids2 = table.search(cr, uid, [(parent, 'in', ids)], context=context)
                    return ids + rg(ids2, table, parent)
                return [(left, 'in', rg(ids, table, parent or table._parent_name))]

        self.__main_table = table

        i = -1
        while i + 1<len(self.__exp):
            i += 1
            e = self.__exp[i]
            if self._is_operator(e) or e == self.__DUMMY_LEAF:
                continue
            left, operator, right = e

            working_table = table
            main_table = table
            fargs = left.split('.', 1)
            index = i
            if left in table._inherit_fields:
                while True:
                    field = main_table._columns.get(fargs[0], False)
                    if field:
                        working_table = main_table
                        self.__tables[i] = working_table
                        break
                    working_table = main_table.pool.get(main_table._inherit_fields[left][0])
                    if working_table not in self.__tables.values():
                        self.__joins.append(('%s.%s=%s.%s' % (working_table._table, 'id', main_table._table, main_table._inherits[working_table._name]), working_table._table))
                        self.__tables[index] = working_table
                        index += 1
                    main_table = working_table

            field = working_table._columns.get(fargs[0], False)
            if not field:
                if left == 'id' and operator == 'child_of':
                    dom = _rec_get(right, working_table)
                    self.__exp = self.__exp[:i] + dom + self.__exp[i+1:]
                continue

            field_obj = table.pool.get(field._obj)
            if len(fargs) > 1:
                if field._type == 'many2one':
                    right = field_obj.search(cr, uid, [(fargs[1], operator, right)], context=context)
                    self.__exp[i] = (fargs[0], 'in', right)
                continue

            if field._properties:
                # this is a function field
                if not field.store:
                    if not field._fnct_search:
                        # the function field doesn't provide a search function and doesn't store
                        # values in the database, so we must ignore it : we generate a dummy leaf
                        self.__exp[i] = self.__DUMMY_LEAF
                    else:
                        subexp = field.search(cr, uid, table, left, [self.__exp[i]])
                        # we assume that the expression is valid
                        # we create a dummy leaf for forcing the parsing of the resulting expression
                        self.__exp[i] = '&'
                        self.__exp.insert(i + 1, self.__DUMMY_LEAF)
                        for j, se in enumerate(subexp):
                            self.__exp.insert(i + 2 + j, se)

                # else, the value of the field is store in the database, so we search on it

            elif field._type == 'one2many':
                if isinstance(right, basestring):
                    ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], operator, limit=None)]
                else:
                    ids2 = list(right)
                if not ids2:
                    self.__exp[i] = ('id', '=', '0')
                else:
                    self.__exp[i] = ('id', 'in', self.__execute_recursive_in(cr, field._fields_id, field_obj._table, 'id', ids2))

            elif field._type == 'many2many':
                #FIXME
                if operator == 'child_of':
                    if isinstance(right, basestring):
                        ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], 'like', limit=None)]
                    else:
                        ids2 = list(right)

                    def _rec_convert(ids):
                        if field_obj == table:
                            return ids
                        return self.__execute_recursive_in(cr, field._id1, field._rel, field._id2, ids)

                    dom = _rec_get(ids2, field_obj)
                    ids2 = field_obj.search(cr, uid, dom, context=context)
                    self.__exp[i] = ('id', 'in', _rec_convert(ids2))
                else:
                    if isinstance(right, basestring):
                        res_ids = [x[0] for x in field_obj.name_search(cr, uid, right, [], operator)]
                    else:
                        res_ids = list(right)
                    self.__exp[i] = ('id', 'in', self.__execute_recursive_in(cr, field._id1, field._rel, field._id2, res_ids) or [0])
            elif field._type == 'many2one':
                if operator == 'child_of':
                    if isinstance(right, basestring):
                        ids2 = [x[0] for x in field_obj.name_search(cr, uid, right, [], 'like', limit=None)]
                    else:
                        ids2 = list(right)

                    self.__operator = 'in'
                    if field._obj != working_table._name:
                        dom = _rec_get(ids2, field_obj, left=left, prefix=field._obj)
                    else:
                        dom = _rec_get(ids2, working_table, parent=left)
                    self.__exp = self.__exp[:i] + dom + self.__exp[i+1:]
                else:
                    if isinstance(right, basestring): # and not isinstance(field, fields.related):
                        c = context.copy()
                        c['active_test'] = False
                        res_ids = field_obj.name_search(cr, uid, right, [], operator, limit=None, context=c)
                        right = map(lambda x: x[0], res_ids)
                        self.__exp[i] = (left, 'in', right)
            else:
                # other field type
                # add the time part to datetime field when it's not there:
                if field._type == 'datetime' and self.__exp[i][2] and len(self.__exp[i][2]) == 10:

                    self.__exp[i] = list(self.__exp[i])

                    if operator in ('>', '>='):
                        self.__exp[i][2] += ' 00:00:00'
                    elif operator in ('<', '<='):
                        self.__exp[i][2] += ' 23:59:59'

                    self.__exp[i] = tuple(self.__exp[i])

                if field.translate:
                    if operator in ('like', 'ilike', 'not like', 'not ilike'):
                        right = '%%%s%%' % right

                    operator = operator == '=like' and 'like' or operator

                    query1 = '( SELECT res_id'          \
                             '    FROM ir_translation'  \
                             '   WHERE name = %s'       \
                             '     AND lang = %s'       \
                             '     AND type = %s'
                    instr = ' %s'
                    #Covering in,not in operators with operands (%s,%s) ,etc.
                    if operator in ['in','not in']:
                        instr = ','.join(['%s'] * len(right))
                        query1 += '     AND value ' + operator +  ' ' +" (" + instr + ")"   \
                             ') UNION ('                \
                             '  SELECT id'              \
                             '    FROM "' + working_table._table + '"'       \
                             '   WHERE "' + left + '" ' + operator + ' ' +" (" + instr + "))"
                    else:
                        query1 += '     AND value ' + operator + instr +   \
                             ') UNION ('                \
                             '  SELECT id'              \
                             '    FROM "' + working_table._table + '"'       \
                             '   WHERE "' + left + '" ' + operator + instr + ")"

                    query2 = [working_table._name + ',' + left,
                              context.get('lang', False) or 'en_US',
                              'model',
                              right,
                              right,
                             ]

                    self.__exp[i] = ('id', 'inselect', (query1, query2))

        return self

    def __leaf_to_sql(self, leaf, table):
        if leaf == self.__DUMMY_LEAF:
            return ('(1=1)', [])
        left, operator, right = leaf

        if operator == 'inselect':
            query = '(%s.%s in (%s))' % (table._table, left, right[0])
            params = right[1]
        elif operator in ['in', 'not in']:
            params = right[:]
            len_before = len(params)
            for i in range(len_before)[::-1]:
                if params[i] == False:
                    del params[i]

            len_after = len(params)
            check_nulls = len_after != len_before
            query = '(1=0)'

            if len_after:
                if left == 'id':
                    instr = ','.join(['%s'] * len_after)
                else:
                    instr = ','.join([table._columns[left]._symbol_set[0]] * len_after)
                query = '(%s.%s %s (%s))' % (table._table, left, operator, instr)

            if check_nulls:
                query = '(%s OR %s.%s IS NULL)' % (query, table._table, left)
        else:
            params = []

            if right == False and (leaf[0] in table._columns)  and table._columns[leaf[0]]._type=="boolean"  and (operator == '='):
                query = '(%s.%s IS NULL or %s.%s = false )' % (table._table, left,table._table, left)
            elif (((right == False) and (type(right)==bool)) or (right is None)) and (operator == '='):
                query = '%s.%s IS NULL ' % (table._table, left)
            elif right == False and (leaf[0] in table._columns)  and table._columns[leaf[0]]._type=="boolean"  and (operator in ['<>', '!=']):
                query = '(%s.%s IS NOT NULL and %s.%s != false)' % (table._table, left,table._table, left)
            elif (((right == False) and (type(right)==bool)) or right is None) and (operator in ['<>', '!=']):
                query = '%s.%s IS NOT NULL' % (table._table, left)
            else:
                if left == 'id':
                    query = '%s.id %s %%s' % (table._table, operator)
                    params = right
                else:
                    like = operator in ('like', 'ilike', 'not like', 'not ilike')

                    op = operator == '=like' and 'like' or operator
                    if left in table._columns:
                        format = like and '%s' or table._columns[left]._symbol_set[0]
                        query = '(%s.%s %s %s)' % (table._table, left, op, format)
                    else:
                        query = "(%s.%s %s '%s')" % (table._table, left, op, right)

                    add_null = False
                    if like:
                        if isinstance(right, str):
                            str_utf8 = right
                        elif isinstance(right, unicode):
                            str_utf8 = right.encode('utf-8')
                        else:
                            str_utf8 = str(right)
                        params = '%%%s%%' % str_utf8
                        add_null = not str_utf8
                    elif left in table._columns:
                        params = table._columns[left]._symbol_set[1](right)

                    if add_null:
                        query = '(%s OR %s IS NULL)' % (query, left)

        if isinstance(params, basestring):
            params = [params]
        return (query, params)

    def to_sql(self):
        stack = []
        params = []
        for i, e in reverse_enumerate(self.__exp):
            if self._is_leaf(e, internal=True):
                table = self.__tables.get(i, self.__main_table)
                q, p = self.__leaf_to_sql(e, table)
                params.insert(0, p)
                stack.append(q)
            else:
                if e == '!':
                    stack.append('(NOT (%s))' % (stack.pop(),))
                else:
                    ops = {'&': ' AND ', '|': ' OR '}
                    q1 = stack.pop()
                    q2 = stack.pop()
                    stack.append('(%s %s %s)' % (q1, ops[e], q2,))

        query = ' AND '.join(reversed(stack))
        joins = ' AND '.join(map(lambda j: j[0], self.__joins))
        if joins:
            query = '(%s) AND (%s)' % (joins, query)
        return (query, flatten(params))

    def get_tables(self):
        return ['"%s"' % t._table for t in set(self.__tables.values()+[self.__main_table])]

# vim:expandtab:smartindent:tabstop=4:softtabstop=4:shiftwidth=4:

<|endoftext|>


<commit_before>import unittest

from apel.db.records import StorageRecord

class StorageRecordTest(unittest.TestCase):
    """
    Test case for StorageRecord
    """

    def test_load_from_tuple(self):
        # full example
        data = ('host.example.org/sr/87912469269276', 1289293612,
                'host.example.org', 'MySite', 'pool-003', 'disk', 'replicated',
                42, '/home/projectA', 'johndoe', 'projectA',
                '/O=Grid/OU=example.org/CN=John Doe',
                'binarydataproject.example.org', 'uk', 'poweruser',
                '2010-10-11T09:31:40Z', '2010-10-11T09:41:40Z', 14728, 13617, 0)

        record = StorageRecord()
        record.load_from_tuple(data)

        self.assertEquals(record.get_field('RecordId'),
                          'host.example.org/sr/87912469269276')
        self.assertEquals(str(record.get_field('CreateTime')),
                          "2010-11-09 09:06:52")
        self.assertEquals(record.get_field('StorageSystem'), 'host.example.org')
        self.assertEquals(record.get_field('Site'), 'MySite')
        self.assertEquals(record.get_field('StorageShare'), 'pool-003')
        self.assertEquals(record.get_field('StorageMedia'), 'disk')
        self.assertEquals(record.get_field('StorageClass'), 'replicated')
        self.assertEquals(record.get_field('FileCount'), 42)
        self.assertEquals(record.get_field('DirectoryPath'), '/home/projectA')
        self.assertEquals(record.get_field('LocalUser'), 'johndoe')
        self.assertEquals(record.get_field('LocalGroup'), 'projectA')
        self.assertEquals(record.get_field('UserIdentity'),
                          '/O=Grid/OU=example.org/CN=John Doe')
        self.assertEquals(record.get_field('Group'),
                          'binarydataproject.example.org')
        self.assertEquals(record.get_field('SubGroup'), 'uk')
        self.assertEquals(record.get_field('Role'), 'poweruser')
        self.assertEquals(str(record.get_field('StartTime')),
                          "2010-10-11 09:31:40")
        self.assertEquals(str(record.get_field('EndTime')),
                          "2010-10-11 09:41:40")
        self.assertEquals(record.get_field('ResourceCapacityUsed'), 14728)
        self.assertEquals(record.get_field('LogicalCapacityUsed'), 13617)
        self.assertEquals(record.get_field('ResourceCapacityAllocated'), 0)

    def test_mandatory_fields(self):
        record = StorageRecord()
        record.set_field('RecordId', 'host.example.org/sr/87912469269276')
        record.set_field('CreateTime', 1289293612)
        record.set_field('StorageSystem', 'host.example.org')
        record.set_field('StartTime', 1286785900)
        record.set_field('EndTime', 1286786500)
        record.set_field('ResourceCapacityUsed', 14728)

        try:
            record._check_fields()
        except Exception, e:
            self.fail('_check_fields method failed: %s [%s]' % (str(e), str(type(e))))

if __name__ == '__main__':
    unittest.main()
<commit_msg>Add test for storage record get_apel_db_insert<commit_after>import unittest

from apel.db.records import StorageRecord

class StorageRecordTest(unittest.TestCase):
    """
    Test case for StorageRecord
    """

    def test_load_from_tuple(self):
        # full example
        data = ('host.example.org/sr/87912469269276', 1289293612,
                'host.example.org', 'MySite', 'pool-003', 'disk', 'replicated',
                42, '/home/projectA', 'johndoe', 'projectA',
                '/O=Grid/OU=example.org/CN=John Doe',
                'binarydataproject.example.org', 'uk', 'poweruser',
                '2010-10-11T09:31:40Z', '2010-10-11T09:41:40Z', 14728, 13617, 0)

        record = StorageRecord()
        record.load_from_tuple(data)

        self.assertEquals(record.get_field('RecordId'),
                          'host.example.org/sr/87912469269276')
        self.assertEquals(str(record.get_field('CreateTime')),
                          "2010-11-09 09:06:52")
        self.assertEquals(record.get_field('StorageSystem'), 'host.example.org')
        self.assertEquals(record.get_field('Site'), 'MySite')
        self.assertEquals(record.get_field('StorageShare'), 'pool-003')
        self.assertEquals(record.get_field('StorageMedia'), 'disk')
        self.assertEquals(record.get_field('StorageClass'), 'replicated')
        self.assertEquals(record.get_field('FileCount'), 42)
        self.assertEquals(record.get_field('DirectoryPath'), '/home/projectA')
        self.assertEquals(record.get_field('LocalUser'), 'johndoe')
        self.assertEquals(record.get_field('LocalGroup'), 'projectA')
        self.assertEquals(record.get_field('UserIdentity'),
                          '/O=Grid/OU=example.org/CN=John Doe')
        self.assertEquals(record.get_field('Group'),
                          'binarydataproject.example.org')
        self.assertEquals(record.get_field('SubGroup'), 'uk')
        self.assertEquals(record.get_field('Role'), 'poweruser')
        self.assertEquals(str(record.get_field('StartTime')),
                          "2010-10-11 09:31:40")
        self.assertEquals(str(record.get_field('EndTime')),
                          "2010-10-11 09:41:40")
        self.assertEquals(record.get_field('ResourceCapacityUsed'), 14728)
        self.assertEquals(record.get_field('LogicalCapacityUsed'), 13617)
        self.assertEquals(record.get_field('ResourceCapacityAllocated'), 0)

    def test_mandatory_fields(self):
        record = StorageRecord()
        record.set_field('RecordId', 'host.example.org/sr/87912469269276')
        record.set_field('CreateTime', 1289293612)
        record.set_field('StorageSystem', 'host.example.org')
        record.set_field('StartTime', 1286785900)
        record.set_field('EndTime', 1286786500)
        record.set_field('ResourceCapacityUsed', 14728)

        try:
            record._check_fields()
        except Exception, e:
            self.fail('_check_fields method failed: %s [%s]' % (str(e), str(type(e))))

    def test_get_apel_db_insert(self):
        """Check that calling get_apel_db_insert raises no exceptions

        Previously there was a mismatch between the number of arguments of this
        and the parent Record class which led to a TypeError being raised.
        """

        record = StorageRecord()
        record.set_field('RecordId', 'host.example.org/sr/87912469269276')
        record.set_field('CreateTime', 1289293612)
        record.set_field('StorageSystem', 'host.example.org')
        record.set_field('StartTime', 1286785900)
        record.set_field('EndTime', 1286786500)
        record.set_field('ResourceCapacityUsed', 14728)

        record.get_apel_db_insert()
        record.get_apel_db_insert(None)

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


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Specify and constraints to determine which targets are observable for
an observer.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from abc import ABCMeta, abstractmethod
import numpy as np
from astropy.time import Time
import astropy.units as u
from astropy.coordinates import get_sun

DEFAULT_TIME_RESOLUTION = 0.5*u.hour

__all__ = ["AltitudeConstraint", "AirmassConstraint", "AtNight",
           "is_observable", "is_always_observable", "time_grid_from_range"]

@u.quantity_input(time_resolution=u.hour)
def time_grid_from_range(time_range, time_resolution=DEFAULT_TIME_RESOLUTION):
    """
    Get linearly-spaced sequence of times.

    Parameters
    ----------
    time_range : `~astropy.time.Time` of length=2
        Lower and upper bounds on time sequence

    time_resolution : `~astropy.units.quantity` (optional)
        Time-grid spacing

    Returns
    -------
    times : `~astropy.time.Time`
        Linearly-spaced sequence of times
    """
    return Time(np.arange(time_range[0].jd, time_range[1].jd,
                          time_resolution.to(u.day).value), format='jd')

class Constraint(object):
    """
    Abstract class for objects defining observational constraints.
    """
    __metaclass__ = ABCMeta

    def __call__(self, time_range, observer, targets):
        cons = self._compute_constraint(time_range, observer, targets)
        return cons

    def _compute_constraint(self, time_range, observer, targets):
        # Should be implemented on each subclass of Constraint
        raise NotImplementedError

    def _get_altaz(self, time_range, observer, targets,
                   time_resolution=DEFAULT_TIME_RESOLUTION,
                   force_zero_pressure=False):
        """
        Calculate alt/az for ``target`` at times linearly spaced between
        the two times in ``time_range`` with grid spacing ``time_resolution``
        for ``observer``.

        Cache the result on the ``observer`` object.

        Parameters
        ----------
        time_range : `~astropy.time.Time` with length=2
            Time range on which to compute these constraints

        targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
            Target or list of targets

        observer : `~astroplan.Observer`
            The observer who has constraints ``constraints``

        time_resolution : `~astropy.units.Quantity` (optional)
            Set the time resolution in calculations of the altitude/azimuth

        Returns
        -------
        altaz_dict : dict
            Dictionary containing two key-value pairs. (1) 'times' contains the
            times for the alt/az computations, (2) 'altaz' contains the
            corresponding alt/az coordinates at those times.
        """
        if not hasattr(observer, '_altaz_cache'):
            observer._altaz_cache = {}

        times = time_grid_from_range(time_range)

        # convert times, targets to tuple for hashing
        aakey = (tuple(times.jd), tuple(targets))

        if aakey not in observer._altaz_cache:
            if force_zero_pressure:
                observer_old_pressure = observer.pressure
                observer.pressure = 0

            altaz = observer.altaz(times, targets)
            observer._altaz_cache[aakey] = dict(times=times,
                                                altaz=altaz)
            if force_zero_pressure:
                observer.pressure = observer_old_pressure

        return observer._altaz_cache[aakey]

class AltitudeConstraint(Constraint):
    """
    Constrain the altitude of the target.

    Note: this will misbehave if you try to constrain negative altitudes, as
    the `~astropy.coordinates.AltAz` frame tends to mishandle negative
    altitudes.
    """
    def __init__(self, min=None, max=None):
        """
        Parameters
        ----------
        min : `~astropy.units.Quantity` or `None`
            Minimum altitude of the target. `None` indicates no limit.

        max : `~astropy.units.Quantity` or `None`
            Maximum altitude of the target. `None` indicates no limit.
        """
        if min is None:
            self.min = 0*u.deg
        else:
            self.min = min
        if max is None:
            self.max = 90*u.deg
        else:
            self.max = max

    def _compute_constraint(self, time_range, observer, targets,
                            time_resolution=DEFAULT_TIME_RESOLUTION):

        cached_altaz = self._get_altaz(time_range, observer, targets)
        altaz = cached_altaz['altaz']
        lowermask = self.min < altaz.alt
        uppermask = altaz.alt < self.max
        return lowermask & uppermask

class AirmassConstraint(AltitudeConstraint):
    """
    Constrain the airmass of a target.

    The airmass is approximated by the secant of the zenith angle.
    """
    def __init__(self, max=None, min=None):
        """
        Note: the ``max`` and ``min`` arguments appear in the order (max, min)
        in this initializer to support the common case for users who care about
        the upper limit on the airmass (``max``) and not the lower limit.

        Parameters
        ----------
        max : float or `None`
            Maximum airmass of the target. `None` indicates no limit.

        min : float or `None`
            Minimum airmass of the target. `None` indicates no limit.
        """
        self.min = min
        self.max = max

    def _compute_constraint(self, time_range, observer, targets):
        cached_altaz = self._get_altaz(time_range, observer, targets)
        altaz = cached_altaz['altaz']
        if self.min is None:
            mask = altaz.secz < self.max
        elif self.max is None:
            mask = self.min < altaz.secz
        elif self.min is not None and self.max is not None:
            mask = (self.min < altaz.secz) & (altaz.secz < self.max)
        return mask

class AtNight(Constraint):
    """
    Constrain the Sun to be below ``horizon``.
    """
    @u.quantity_input(horizon=u.deg)
    def __init__(self, max_solar_altitude=0*u.deg, force_pressure_zero=True):
        """
        Parameters
        ----------
        max_solar_altitude : `~astropy.units.Quantity`
            Define "night" as when the sun is below ``max_solar_altitude``.
            Default is zero degrees altitude.
        """
        self.max_solar_altitude = max_solar_altitude
        self.force_pressure_zero = force_pressure_zero

    @classmethod
    def twilight_civil(cls, **kwargs):
        return cls(max_solar_altitude=-6*u.deg, **kwargs)

    @classmethod
    def twilight_nautical(cls, **kwargs):
        return cls(max_solar_altitude=-12*u.deg, **kwargs)

    @classmethod
    def twilight_astronomical(cls, **kwargs):
        return cls(max_solar_altitude=-18*u.deg, **kwargs)

    def _get_solar_altitudes(self, time_range, observer, targets,
                             time_resolution=DEFAULT_TIME_RESOLUTION):
        if not hasattr(observer, '_altaz_cache'):
            observer._altaz_cache = {}

        times = Time(np.arange(time_range[0].jd, time_range[1].jd,
                               time_resolution.to(u.day).value), format='jd')

        aakey = (tuple(times.jd), 'sun')

        if aakey not in observer._altaz_cache:
            if self.force_pressure_zero:
                observer_old_pressure = observer.pressure
                observer.pressure = 0

            # Broadcast the solar altitudes for the number of targets:
            altaz = observer.altaz(times, get_sun(times))
            altitude = altaz.alt
            altitude.resize(1, len(altitude))
            altitude = altitude + np.zeros((len(targets), 1))

            observer._altaz_cache[aakey] = dict(times=times,
                                                altitude=altitude)
            if self.force_pressure_zero:
                observer.pressure = observer_old_pressure

        return observer._altaz_cache[aakey]

    def _compute_constraint(self, time_range, observer, targets):
        sun_altaz = self._get_solar_altitudes(time_range, observer, targets)
        solar_altitude = sun_altaz['altitude']
        mask = solar_altitude < self.max_solar_altitude
        return mask

def is_always_observable(constraints, time_range, targets, observer):
    """
    Are the ``targets`` always observable throughout ``time_range`` given
    constraints in ``constraints_list`` for ``observer``?

    Parameters
    ----------
    constraints : list or `~astroplan.constraints.Constraint`
        Observational constraint(s)

    time_range : `~astropy.time.Time` with length=2
        Time range on which to compute these constraints

    targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
        Target or list of targets

    observer : `~astroplan.Observer`
        The observer who has constraints ``constraints``

    Returns
    -------
    ever_observable : list
        List of booleans of same length as ``targets`` for whether or not each
        target is observable in the time range given the constraints.
    """
    if not hasattr(constraints, '__len__'):
        constraints = [constraints]

    contraint_arr = np.logical_and.reduce([constraint(time_range, observer,
                                                      targets)
                                          for constraint in constraints])
    return np.all(contraint_arr, axis=1)

def is_observable(constraints, time_range, targets, observer):
    """
    Are the ``targets`` observable during ``time_range`` given constraints in
    ``constraints_list`` for ``observer``?

    Parameters
    ----------
    constraints : list or `~astroplan.constraints.Constraint`
        Observational constraint(s)

    time_range : `~astropy.time.Time` with length=2
        Time range on which to compute these constraints

    targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
        Target or list of targets

    observer : `~astroplan.Observer`
        The observer who has constraints ``constraints``

    Returns
    -------
    ever_observable : list
        List of booleans of same length as ``targets`` for whether or not each
        target is ever observable in the time range given the constraints.
    """
    if not hasattr(constraints, '__len__'):
        constraints = [constraints]

    contraint_arr = np.logical_and.reduce([constraint(time_range, observer,
                                                      targets)
                                          for constraint in constraints])
    return np.any(contraint_arr, axis=1)
<commit_msg>Adding a SunSeparation constraint<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Specify and constraints to determine which targets are observable for
an observer.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from abc import ABCMeta, abstractmethod
import numpy as np
from astropy.time import Time
import astropy.units as u
from astropy.coordinates import get_sun, Angle

DEFAULT_TIME_RESOLUTION = 0.5*u.hour

__all__ = ["AltitudeConstraint", "AirmassConstraint", "AtNight",
           "is_observable", "is_always_observable", "time_grid_from_range"]

@u.quantity_input(time_resolution=u.hour)
def time_grid_from_range(time_range, time_resolution=DEFAULT_TIME_RESOLUTION):
    """
    Get linearly-spaced sequence of times.

    Parameters
    ----------
    time_range : `~astropy.time.Time` of length=2
        Lower and upper bounds on time sequence

    time_resolution : `~astropy.units.quantity` (optional)
        Time-grid spacing

    Returns
    -------
    times : `~astropy.time.Time`
        Linearly-spaced sequence of times
    """
    return Time(np.arange(time_range[0].jd, time_range[1].jd,
                          time_resolution.to(u.day).value), format='jd')

class Constraint(object):
    """
    Abstract class for objects defining observational constraints.
    """
    __metaclass__ = ABCMeta

    def __call__(self, time_range, observer, targets):
        cons = self._compute_constraint(time_range, observer, targets)
        return cons

    def _compute_constraint(self, time_range, observer, targets):
        # Should be implemented on each subclass of Constraint
        raise NotImplementedError

    def _get_altaz(self, time_range, observer, targets,
                   time_resolution=DEFAULT_TIME_RESOLUTION,
                   force_zero_pressure=False):
        """
        Calculate alt/az for ``target`` at times linearly spaced between
        the two times in ``time_range`` with grid spacing ``time_resolution``
        for ``observer``.

        Cache the result on the ``observer`` object.

        Parameters
        ----------
        time_range : `~astropy.time.Time` with length=2
            Time range on which to compute these constraints

        targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
            Target or list of targets

        observer : `~astroplan.Observer`
            The observer who has constraints ``constraints``

        time_resolution : `~astropy.units.Quantity` (optional)
            Set the time resolution in calculations of the altitude/azimuth

        Returns
        -------
        altaz_dict : dict
            Dictionary containing two key-value pairs. (1) 'times' contains the
            times for the alt/az computations, (2) 'altaz' contains the
            corresponding alt/az coordinates at those times.
        """
        if not hasattr(observer, '_altaz_cache'):
            observer._altaz_cache = {}

        times = time_grid_from_range(time_range)

        # convert times, targets to tuple for hashing
        aakey = (tuple(times.jd), tuple(targets))

        if aakey not in observer._altaz_cache:
            if force_zero_pressure:
                observer_old_pressure = observer.pressure
                observer.pressure = 0

            altaz = observer.altaz(times, targets)
            observer._altaz_cache[aakey] = dict(times=times,
                                                altaz=altaz)
            if force_zero_pressure:
                observer.pressure = observer_old_pressure

        return observer._altaz_cache[aakey]

class AltitudeConstraint(Constraint):
    """
    Constrain the altitude of the target.

    Note: this will misbehave if you try to constrain negative altitudes, as
    the `~astropy.coordinates.AltAz` frame tends to mishandle negative
    altitudes.
    """
    def __init__(self, min=None, max=None):
        """
        Parameters
        ----------
        min : `~astropy.units.Quantity` or `None`
            Minimum altitude of the target. `None` indicates no limit.

        max : `~astropy.units.Quantity` or `None`
            Maximum altitude of the target. `None` indicates no limit.
        """
        if min is None:
            self.min = 0*u.deg
        else:
            self.min = min
        if max is None:
            self.max = 90*u.deg
        else:
            self.max = max

    def _compute_constraint(self, time_range, observer, targets,
                            time_resolution=DEFAULT_TIME_RESOLUTION):

        cached_altaz = self._get_altaz(time_range, observer, targets)
        altaz = cached_altaz['altaz']
        lowermask = self.min < altaz.alt
        uppermask = altaz.alt < self.max
        return lowermask & uppermask

class AirmassConstraint(AltitudeConstraint):
    """
    Constrain the airmass of a target.

    The airmass is approximated by the secant of the zenith angle.
    """
    def __init__(self, max=None, min=None):
        """
        Note: the ``max`` and ``min`` arguments appear in the order (max, min)
        in this initializer to support the common case for users who care about
        the upper limit on the airmass (``max``) and not the lower limit.

        Parameters
        ----------
        max : float or `None`
            Maximum airmass of the target. `None` indicates no limit.

        min : float or `None`
            Minimum airmass of the target. `None` indicates no limit.
        """
        self.min = min
        self.max = max

    def _compute_constraint(self, time_range, observer, targets):
        cached_altaz = self._get_altaz(time_range, observer, targets)
        altaz = cached_altaz['altaz']
        if self.min is None and self.max is not None:
            mask = altaz.secz < self.max
        elif self.max is None and self.min is not None:
            mask = self.min < altaz.secz
        elif self.min is not None and self.max is not None:
            mask = (self.min < altaz.secz) & (altaz.secz < self.max)
        else:
            raise ValueError("No max and/or min specified in "
                             "AirmassConstraint.")
        return mask

class AtNight(Constraint):
    """
    Constrain the Sun to be below ``horizon``.
    """
    @u.quantity_input(horizon=u.deg)
    def __init__(self, max_solar_altitude=0*u.deg, force_pressure_zero=True):
        """
        Parameters
        ----------
        max_solar_altitude : `~astropy.units.Quantity`
            Define "night" as when the sun is below ``max_solar_altitude``.
            Default is zero degrees altitude.
        """
        self.max_solar_altitude = max_solar_altitude
        self.force_pressure_zero = force_pressure_zero

    @classmethod
    def twilight_civil(cls, **kwargs):
        return cls(max_solar_altitude=-6*u.deg, **kwargs)

    @classmethod
    def twilight_nautical(cls, **kwargs):
        return cls(max_solar_altitude=-12*u.deg, **kwargs)

    @classmethod
    def twilight_astronomical(cls, **kwargs):
        return cls(max_solar_altitude=-18*u.deg, **kwargs)

    def _get_solar_altitudes(self, time_range, observer, targets,
                             time_resolution=DEFAULT_TIME_RESOLUTION):
        if not hasattr(observer, '_altaz_cache'):
            observer._altaz_cache = {}

        times = time_grid_from_range(time_range)

        aakey = (tuple(times.jd), 'sun')

        if aakey not in observer._altaz_cache:
            if self.force_pressure_zero:
                observer_old_pressure = observer.pressure
                observer.pressure = 0

            # Broadcast the solar altitudes for the number of targets:
            altaz = observer.altaz(times, get_sun(times))
            altitude = altaz.alt
            altitude.resize(1, len(altitude))
            altitude = altitude + np.zeros((len(targets), 1))

            observer._altaz_cache[aakey] = dict(times=times,
                                                altitude=altitude)
            if self.force_pressure_zero:
                observer.pressure = observer_old_pressure

        return observer._altaz_cache[aakey]

    def _compute_constraint(self, time_range, observer, targets):
        sun_altaz = self._get_solar_altitudes(time_range, observer, targets)
        solar_altitude = sun_altaz['altitude']
        mask = solar_altitude < self.max_solar_altitude
        return mask

class SunSeparation(Constraint):
    """
    Constrain the sun to be a distance away from target.
    """
    def __init__(self, min=None, max=None):
        self.min = min
        self.max = max

    def _compute_constraint(self, time_range, observer, targets):
        times = time_grid_from_range(time_range)
        sun = get_sun(times)
        targets = [target.coord if hasattr(target, 'coord') else target
                   for target in targets]
        solar_separation = Angle([sun.separation(target) for target in targets])
        if self.min is None and self.max is not None:
            mask = self.max > solar_separation
        elif self.max is None and self.min is not None:
            mask = self.min < solar_separation
        elif self.min is not None and self.max is not None:
            mask = ((self.min < solar_separation) &
                    (solar_separation < self.max))
        else:
            raise ValueError("No max and/or min specified in "
                             "SunSeparation.")
        return mask

def is_always_observable(constraints, time_range, targets, observer):
    """
    Are the ``targets`` always observable throughout ``time_range`` given
    constraints in ``constraints_list`` for ``observer``?

    Parameters
    ----------
    constraints : list or `~astroplan.constraints.Constraint`
        Observational constraint(s)

    time_range : `~astropy.time.Time` with length=2
        Time range on which to compute these constraints

    targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
        Target or list of targets

    observer : `~astroplan.Observer`
        The observer who has constraints ``constraints``

    Returns
    -------
    ever_observable : list
        List of booleans of same length as ``targets`` for whether or not each
        target is observable in the time range given the constraints.
    """
    if not hasattr(constraints, '__len__'):
        constraints = [constraints]

    contraint_arr = np.logical_and.reduce([constraint(time_range, observer,
                                                      targets)
                                          for constraint in constraints])
    return np.all(contraint_arr, axis=1)

def is_observable(constraints, time_range, targets, observer):
    """
    Are the ``targets`` observable during ``time_range`` given constraints in
    ``constraints_list`` for ``observer``?

    Parameters
    ----------
    constraints : list or `~astroplan.constraints.Constraint`
        Observational constraint(s)

    time_range : `~astropy.time.Time` with length=2
        Time range on which to compute these constraints

    targets : {list, `~astropy.coordinates.SkyCoord`, `~astroplan.FixedTarget`}
        Target or list of targets

    observer : `~astroplan.Observer`
        The observer who has constraints ``constraints``

    Returns
    -------
    ever_observable : list
        List of booleans of same length as ``targets`` for whether or not each
        target is ever observable in the time range given the constraints.
    """
    if not hasattr(constraints, '__len__'):
        constraints = [constraints]

    contraint_arr = np.logical_and.reduce([constraint(time_range, observer,
                                                      targets)
                                          for constraint in constraints])
    return np.any(contraint_arr, axis=1)
<|endoftext|>


<commit_before>import InstrumentationManifestTools.manifest as manifest
import InstrumentationManifestTools.manifest_generator as manifest_generator
import InstrumentationManifestTools.wprp_generator as wprp_generator

def sample():
    p = manifest.Provider("Multi-Main", binary_filename = "%temp%/test.dll")

    task = manifest.Task("task1")
    p.add(task)

    op = manifest.Opcode("opcode")
    p.add(op)

    kw = manifest.Keyword("kw", mask = '0x1')
    p.add(kw)

    filter = manifest.Filter("fil")
    p.add(filter)

    lev = manifest.Level("lev")
    p.add(lev)

    templ = manifest.Template("cha")
    templ.add_data("Description", "win:AnsiString")
    p.add(templ)

    channel = manifest.Channel("kchaw")
    p.add(channel)

    ev = manifest.Event("pork",
        channel = channel,
        task = task,
        opcode = op,
        keywords = kw,
        level = lev,
        template = templ
        )
    p.add(ev)

    ev = manifest.Event("pork2",
        template = templ
        )
    p.add(ev)

    with open('sample.man', 'w') as file:
        file.write(manifest_generator.to_manifest_xml([p]))

    profile_1 = manifest.Profile("General", "Sweet test profile")
    profile_1.add(p)

    with open('sample.wprp', 'w') as file:
        file.write(wprp_generator.to_wprp_xml([profile_1]))

sample()
<commit_msg>reformat test<commit_after>import InstrumentationManifestTools.manifest as manifest
import InstrumentationManifestTools.manifest_generator as manifest_generator
import InstrumentationManifestTools.wprp_generator as wprp_generator

def sample():
    p = manifest.Provider("Multi-Main", binary_filename="%temp%/test.dll")

    task = manifest.Task("task1")
    p.add(task)

    op = manifest.Opcode("opcode")
    p.add(op)

    kw = manifest.Keyword("kw", mask='0x1')
    p.add(kw)

    filter = manifest.Filter("fil")
    p.add(filter)

    lev = manifest.Level("lev")
    p.add(lev)

    templ = manifest.Template("cha")
    templ.add_data("Description", "win:AnsiString")
    p.add(templ)

    channel = manifest.Channel("kchaw")
    p.add(channel)

    ev = manifest.Event("pork",
                        channel=channel,
                        task=task,
                        opcode=op,
                        keywords=kw,
                        level=lev,
                        template=templ
                        )
    p.add(ev)

    ev = manifest.Event("pork2",
                        template=templ
                        )
    p.add(ev)

    with open('sample.man', 'w') as file:
        file.write(manifest_generator.to_manifest_xml([p]))

    profile_1 = manifest.Profile("General", "Sweet test profile")
    profile_1.add(p)

    with open('sample.wprp', 'w') as file:
        file.write(wprp_generator.to_wprp_xml([profile_1]))

sample()
<|endoftext|>