<commit_before># pdbtools.py
#
# PDB reading and writing functions from packmoltools/resolvate, which provide
# rudimentry input and output for pdbfiles.
#
# written by John Chodera, modified by Vincennt Voelz, Stanford University (c) 2007
# -----------------------------------------------------
#
# -----------------------------------------------------
# TO DO
# - This should be revamped with the mmLib module
# -----------------------------------------------------
# MOD HISTORY
# 07/10/2007 - Created pdbtools.py
#            - added ignh=True flag to writeAtomsToPDB()
#            - enabled 4-letter gmx-style ffamber residue names (for termini e.g.)

# GLOBAL imports

import sys, os, os.path
import commands, shutil, tempfile, math

from mmtools.gromacstools.GromacsData import *

def readAtomsFromPDB(pdbfilename):
    """Read atom records from the PDB and return them in a list.

    present_sequence = getPresentSequence(pdbfilename, chain=' ')

    REQUIRED ARGUMENTS
      pdbfilename - the filename of the PDB file to import from

    OPTIONAL ARGUMENTS
      chain - the one-character chain ID of the chain to import (default ' ')

    RETURN VALUES
      sequence (dictionary) - sequence[residue_id] is the one-letter code corresponding to residue index residue_id

    The ATOM records are read, and the sequence for which there are atomic coordinates is stored.

    """

    # Read the PDB file into memory.
    pdbfile = open(pdbfilename, 'r')
    lines = pdbfile.readlines()
    pdbfile.close()

    # Read atoms.
    atoms = [ ]
    for line in lines:
        if line[0:6] == "ATOM  ":
            # Parse line into fields.
            atom = { }
            atom["serial"] = int(line[6:11])
            atom["name"] = line[12:16]
            atom["altLoc"] = line[16:17]
            atom["resName"] = line[17:21]
            atom["chainID"] = line[21:22]
            atom["resSeq"] = int(line[22:26])
            atom["iCode"] = line[26:27]
            atom["x"] = float(line[30:38])
            atom["y"] = float(line[38:46])
            atom["z"] = float(line[46:54])

            atom["occupancy"] = 1.0
            if (line[54:60].strip() != ''):
              atom["occupancy"] = float(line[54:60])

            atom["tempFactor"] = 0.0
            if (line[60:66].strip() != ''):
              atom["tempFactor"] = float(line[60:66])

            atom["segID"] = line[72:76]
            atom["element"] = line[76:78]
            atom["charge"] = line[78:80]

            atoms.append(atom)

    # Return list of atoms.
    return atoms

def writeAtomsToPDB(pdbfilename, atoms, renumber = False, ignh = False):
    """Write atom records to PDB file.

    REQUIRED ARGUMENTS
      pdbfilename - the name of the PDB file to write
      atoms - a list of atom dictionaries -- see readAtomsFromPdb

    OPTIONAL ARGUMENTS
      if renumber is True, then the atom and residue numbers will be renumbered starting with 1

    RETURN VALUES
      none

    EXAMPLE
    writeAtomsToPdb(pdbfilename, atoms)

    """

    # Renumber if desired.
    if (renumber):
      first_residue = atoms[0]["resSeq"]
      serial = 1
      resSeq = 0
      last_resSeq = None
      for atom in atoms:
        atom["serial"] = serial
        serial += 1

        if(atom["resSeq"] != last_resSeq):
          resSeq += 1
          last_resSeq = atom["resSeq"]
        atom["resSeq"] = resSeq

    # Read the PDB file into memory.
    pdbfile = open(pdbfilename, 'w')

    # Write atoms.
    for atom in atoms:
      if (not ignh) or (not isAtomNameHydrogen(atom["name"] )):   # Ignore hydrogens if desired
        pdbfile.write('ATOM  %(serial)5d %(name)4s%(altLoc)c%(resName)4s%(chainID)c%(resSeq)4d%(iCode)c   %(x)8.3f%(y)8.3f%(z)8.3f%(occupancy)6.2f%(tempFactor)6.2f%(element)2s%(charge)2s\n' % atom)
    pdbfile.close()

def isAtomNameHydrogen( atomName ):
    """Returns True if the atomName denotes a non-water hydrogen."""

    if atomName[1] == 'H':
        if atomName.count('W') == 0:    # water hydrogens don't count
            return True
    return False

def getChargedNames(atoms):
    """Find out which residues are charged based on their AMBER-style residue names.

    REQUIRED ARGUMENTS
    atoms - a list of atom dictionaries - see readAtomsFromPdb

    RETURN VALUES
    chargedRes, a dictionary { resSeq: ('+' or '-' or '0', "resName") } depending on whether the titratible side chain residue charged or netural

    NOTES
    The criteria here come from David Mobley renaming fixes in mmtools/mccetools."""

    chargedRes = {}    # { resSeq: '+' or '-'}
    for atom in atoms:

        # tritratible residues
        if atom["resName"].strip() == 'HIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'HIS':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'HID':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'HIE':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'LYN':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'LYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'CYN':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'CYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])
        elif atom["resName"].strip() == 'ASH':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'GLH':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])

        # remember all 4-letter residue names, which are the termini
        if len(atom["resName"].strip()) == 4:
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])

        # Additonally, remember the charged names
        #
        # N-termini names
        if atom["resName"].strip() == 'NHIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        # gmx ffamber does not have CLYN, only CLYP
        elif atom["resName"].strip() == 'NLYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'NCYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])
        # C-termini names
        elif atom["resName"].strip() == 'CHIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        # gmx ffamber does not have CLYN, only CLYP
        elif atom["resName"].strip() == 'CLYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'CCYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])

    return chargedRes

def setChargedNames(atoms, chargedRes):
    """Change the names of the residues based on the charged state indicated in the chargedRes dict

    REQUIRED ARGUMENTS
    atoms      - a list of atom dictionaries -- see readAtomsFromPdb
    chargedRes - a dictionary { resSeq: '+' or '-' or '0' }

    RETURN VALUES
    atoms      - a list of atom dictionaries, renamed to residues in chargedRes dictionary

    NOTE: This only works if we ignore the hydrogens when writing the PDB."""

    for atom in atoms:
        if chargedRes.has_key(atom["resSeq"] ):
            atom["resName"] = chargedRes[ atom["resSeq"] ][1]

    return atoms

def writeGrofileFromPDB(pdbfile, grofile):
    """Take in a PDB file and write it as a *.gro file."""

    gstruct = GromacsStructureFromPDB(pdbfile)
    gstruct.write(grofile)

def GromacsStructureFromPDB(pdbfile):
    """Take in a PDB file a GromacsStructure class."""
    atoms = readAtomsFromPDB(pdbfile)
    gstruct = GromacsStructure(name=pdbfile, header="title" )
    gatoms = GromacsAtoms()
    for atom in atoms:
        gatoms.append( GromacsAtom(resnum=atom["resSeq"], resname=atom["resName"], atomname=atom["name"], atomnum=atom["serial"],
                                      x=(atom["x"]/10.), y=(atom["y"]/10.), z=(atom["z"]/10.), vx=0.0, vy=0.0, vz=0.0) )
    gstruct.appendatoms( gatoms )
    return gstruct

def GromacsStructureFromGrofile(grofile):

    gstruct = GromacsStructure(name=grofile, header="title" )
    gstruct.load(grofile)
    return gstruct

def writePDBFromGrofile(grofile, pdbfile, stripWaters=False, stripIons=False):
    """Take in a *.gro Grofile file and write it as a PDB file."""

    gstruct = GromacsStructure(name=grofile, header="title" )
    gstruct.load(grofile)
    atoms = []
    for gatom in gstruct.atoms:

        atom = { }
        atom["serial"] = gatom.atomnum
        atom["name"] = '%-4s'%(gatom.atomname.strip())
        atom["altLoc"] = ' '
        atom["resName"] = '%-4s'%(gatom.resname.strip())
        atom["chainID"] = ' '
        atom["resSeq"] = gatom.resnum
        atom["iCode"] = ' '
        atom["x"] = gatom.x*10.    # *.gro files are in nm
        atom["y"] = gatom.y*10.
        atom["z"] = gatom.z*10.
        atom["occupancy"] = 1.0
        atom["tempFactor"] = 0.0
        atom["segID"] = '    '
        atom["element"] = '  '
        atom["charge"] = '   '

        if (atom["resName"].count('HOH')>0) or (atom["resName"].count('SOL')>0):
          if not stripWaters:
            atoms.append(atom)
        elif (atom["resName"].count('Cl')>0) or (atom["resName"].count('Na')>0):
          if not stripIons:
            atoms.append(atom)
        else:
            atoms.append(atom)

    writeAtomsToPDB(pdbfile, atoms, renumber = True)

<commit_msg>Added termini conversion tools.<commit_after># pdbtools.py
#
# PDB reading and writing functions from packmoltools/resolvate, which provide
# rudimentry input and output for pdbfiles.
#
# written by John Chodera
# modified by Vincent Voelz
# Stanford University (c) 2007
# -----------------------------------------------------
#
# -----------------------------------------------------
# TO DO
# - This should be revamped with the mmLib module
# -----------------------------------------------------
# MOD HISTORY
# 07/10/2007 - Created pdbtools.py
#            - added ignh=True flag to writeAtomsToPDB()
#            - enabled 4-letter gmx-style ffamber residue names (for termini e.g.)

# GLOBAL imports

import sys, os, os.path
import commands, shutil, tempfile, math

from mmtools.gromacstools.GromacsData import *

def readAtomsFromPDB(pdbfilename):
    """Read atom records from the PDB and return them in a list.

    present_sequence = getPresentSequence(pdbfilename, chain=' ')

    REQUIRED ARGUMENTS
      pdbfilename - the filename of the PDB file to import from

    OPTIONAL ARGUMENTS
      chain - the one-character chain ID of the chain to import (default ' ')

    RETURN VALUES
      sequence (dictionary) - sequence[residue_id] is the one-letter code corresponding to residue index residue_id

    The ATOM records are read, and the sequence for which there are atomic coordinates is stored.

    """

    # Read the PDB file into memory.
    pdbfile = open(pdbfilename, 'r')
    lines = pdbfile.readlines()
    pdbfile.close()

    # Read atoms.
    atoms = [ ]
    for line in lines:
        if line[0:6] == "ATOM  ":
            # Parse line into fields.
            atom = { }
            atom["serial"] = int(line[6:11])
            atom["name"] = line[12:16]
            atom["altLoc"] = line[16:17]
            atom["resName"] = line[17:21]
            atom["chainID"] = line[21:22]
            atom["resSeq"] = int(line[22:26])
            atom["iCode"] = line[26:27]
            atom["x"] = float(line[30:38])
            atom["y"] = float(line[38:46])
            atom["z"] = float(line[46:54])

            atom["occupancy"] = 1.0
            if (line[54:60].strip() != ''):
              atom["occupancy"] = float(line[54:60])

            atom["tempFactor"] = 0.0
            if (line[60:66].strip() != ''):
              atom["tempFactor"] = float(line[60:66])

            atom["segID"] = line[72:76]
            atom["element"] = line[76:78]
            atom["charge"] = line[78:80]

            atoms.append(atom)

    # Return list of atoms.
    return atoms

def writeAtomsToPDB(pdbfilename, atoms, renumber = False, ignh = False):
    """Write atom records to PDB file.

    REQUIRED ARGUMENTS
      pdbfilename - the name of the PDB file to write
      atoms - a list of atom dictionaries -- see readAtomsFromPdb

    OPTIONAL ARGUMENTS
      if renumber is True, then the atom and residue numbers will be renumbered starting with 1

    RETURN VALUES
      none

    EXAMPLE
    writeAtomsToPdb(pdbfilename, atoms)

    """

    # Renumber if desired.
    if (renumber):
      first_residue = atoms[0]["resSeq"]
      serial = 1
      resSeq = 0
      last_resSeq = None
      for atom in atoms:
        atom["serial"] = serial
        serial += 1

        if(atom["resSeq"] != last_resSeq):
          resSeq += 1
          last_resSeq = atom["resSeq"]
        atom["resSeq"] = resSeq

    # Read the PDB file into memory.
    pdbfile = open(pdbfilename, 'w')

    # Write atoms.
    for atom in atoms:
      if (not ignh) or (not isAtomNameHydrogen(atom["name"] )):   # Ignore hydrogens if desired
        pdbfile.write('ATOM  %(serial)5d %(name)4s%(altLoc)c%(resName)4s%(chainID)c%(resSeq)4d%(iCode)c   %(x)8.3f%(y)8.3f%(z)8.3f%(occupancy)6.2f%(tempFactor)6.2f%(element)2s%(charge)2s\n' % atom)
    pdbfile.close()

def isAtomNameHydrogen( atomName ):
    """Returns True if the atomName denotes a non-water hydrogen."""

    if atomName[1] == 'H':
        if atomName.count('W') == 0:    # water hydrogens don't count
            return True
    return False

def getChargedNames(atoms):
    """Find out which residues are charged based on their AMBER-style residue names.

    REQUIRED ARGUMENTS
    atoms - a list of atom dictionaries - see readAtomsFromPdb

    RETURN VALUES
    chargedRes, a dictionary { resSeq: ('+' or '-' or '0', "resName") } depending on whether the titratible side chain residue charged or netural

    NOTES
    The criteria here come from David Mobley renaming fixes in mmtools/mccetools."""

    chargedRes = {}    # { resSeq: '+' or '-'}
    for atom in atoms:

        # tritratible residues
        if atom["resName"].strip() == 'HIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'HIS':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'HID':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'HIE':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'LYN':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'LYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'CYN':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'CYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])
        elif atom["resName"].strip() == 'ASH':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])
        elif atom["resName"].strip() == 'GLH':
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])

        # remember all 4-letter residue names, which are the termini
        if len(atom["resName"].strip()) == 4:
            chargedRes[ atom["resSeq"] ] = ('0', atom["resName"])

        # Additonally, remember the charged names
        #
        # N-termini names
        if atom["resName"].strip() == 'NHIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        # gmx ffamber does not have CLYN, only CLYP
        elif atom["resName"].strip() == 'NLYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'NCYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])
        # C-termini names
        elif atom["resName"].strip() == 'CHIP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        # gmx ffamber does not have CLYN, only CLYP
        elif atom["resName"].strip() == 'CLYP':
            chargedRes[ atom["resSeq"] ] = ('+', atom["resName"])
        elif atom["resName"].strip() == 'CCYM':
            chargedRes[ atom["resSeq"] ] = ('-', atom["resName"])

    return chargedRes

def setChargedNames(atoms, chargedRes):
    """Change the names of the residues based on the charged state indicated in the chargedRes dict

    REQUIRED ARGUMENTS
    atoms      - a list of atom dictionaries -- see readAtomsFromPdb
    chargedRes - a dictionary { resSeq: '+' or '-' or '0' }

    RETURN VALUES
    atoms      - a list of atom dictionaries, renamed to residues in chargedRes dictionary

    NOTE: This only works if we ignore the hydrogens when writing the PDB."""

    for atom in atoms:
        if chargedRes.has_key(atom["resSeq"] ):
            atom["resName"] = chargedRes[ atom["resSeq"] ][1]

    return atoms

def writeGrofileFromPDB(pdbfile, grofile):
    """Take in a PDB file and write it as a *.gro file."""

    gstruct = GromacsStructureFromPDB(pdbfile)
    gstruct.write(grofile)

def GromacsStructureFromPDB(pdbfile):
    """Take in a PDB file a GromacsStructure class."""
    atoms = readAtomsFromPDB(pdbfile)
    gstruct = GromacsStructure(name=pdbfile, header="title" )
    gatoms = GromacsAtoms()
    for atom in atoms:
        gatoms.append( GromacsAtom(resnum=atom["resSeq"], resname=atom["resName"], atomname=atom["name"], atomnum=atom["serial"],
                                      x=(atom["x"]/10.), y=(atom["y"]/10.), z=(atom["z"]/10.), vx=0.0, vy=0.0, vz=0.0) )
    gstruct.appendatoms( gatoms )
    return gstruct

def GromacsStructureFromGrofile(grofile):

    gstruct = GromacsStructure(name=grofile, header="title" )
    gstruct.load(grofile)
    return gstruct

def writePDBFromGrofile(grofile, pdbfile, stripWaters=False, stripIons=False):
    """Take in a *.gro Grofile file and write it as a PDB file."""

    gstruct = GromacsStructure(name=grofile, header="title" )
    gstruct.load(grofile)
    atoms = []
    for gatom in gstruct.atoms:

        atom = { }
        atom["serial"] = gatom.atomnum
        atom["name"] = '%-4s'%(gatom.atomname.strip())
        atom["altLoc"] = ' '
        atom["resName"] = '%-4s'%(gatom.resname.strip())
        atom["chainID"] = ' '
        atom["resSeq"] = gatom.resnum
        atom["iCode"] = ' '
        atom["x"] = gatom.x*10.    # *.gro files are in nm
        atom["y"] = gatom.y*10.
        atom["z"] = gatom.z*10.
        atom["occupancy"] = 1.0
        atom["tempFactor"] = 0.0
        atom["segID"] = '    '
        atom["element"] = '  '
        atom["charge"] = '   '

        if (atom["resName"].count('HOH')>0) or (atom["resName"].count('SOL')>0):
          if not stripWaters:
            atoms.append(atom)
        elif (atom["resName"].count('Cl')>0) or (atom["resName"].count('Na')>0):
          if not stripIons:
            atoms.append(atom)
        else:
            atoms.append(atom)

    writeAtomsToPDB(pdbfile, atoms, renumber = True)

def rebuildPatchedTermini(outPdbFile):
    """Takes in a PDB file with 'patched' termini from Modeller, and rewrites it with the ffAMBER-style
    ACE and NH2 and residues."""

    # The Modeller PDBs look like this:
    """
     EXPDTA    THEORETICAL MODEL, MODELLER 9v1 2007/08/20 13:48:30
     REMARK   6 MODELLER OBJECTIVE FUNCTION:      4986.1367
     REMARK   6 MODELLER BEST TEMPLATE % SEQ ID: 100.000
     ATOM      1  N   LEU     1      -7.461   4.812   8.221  1.00387.74       1SG   2
     ATOM      2  CA  LEU     1      -7.033   6.174   8.059  1.00387.74       1SG   3
     ATOM      3  CB  LEU     1      -7.447   7.022   9.265  1.00387.74       1SG   4
     ATOM      4  CG  LEU     1      -8.557   8.090   8.962  1.00387.74       1SG   5
     ATOM      5  CD1 LEU     1      -8.754   8.471   7.469  1.00387.74       1SG   6
     ATOM      6  CD2 LEU     1      -9.863   7.852   9.710  1.00387.74       1SG   7
     ATOM      7  C   LEU     1      -5.625   6.415   7.559  1.00387.74       1SG   8
     ATOM      8  O   LEU     1      -4.618   5.810   7.920  1.00387.74       1SG   9
cap* ATOM      9  CAY LEU     1      -9.681   4.051   7.379  1.00387.74       1SG  10
cap* ATOM     10  CY  LEU     1      -8.330   4.670   7.236  1.00387.74       1SG  11
cap* ATOM     11  OY  LEU     1      -8.069   5.220   6.173  1.00387.74       1SG  12
     ATOM     12  N   SER     2      -5.488   7.387   6.667  1.00239.19       1SG  13
     ATOM     13  CA  SER     2      -5.907   8.699   7.035  1.00239.19       1SG  14
     ....
     ATOM    279  C   LEU    34     -14.712  13.829   3.507  1.00350.49       1SG 280
     ATOM    280  O   LEU    34     -13.790  13.944   2.701  1.00350.49       1SG 281
     ATOM    281  N   PHE    35     -14.937  14.531   4.654  1.00213.39       1SG 282
     ATOM    282  CA  PHE    35     -14.196  14.588   5.890  1.00213.39       1SG 283
     ATOM    283  CB  PHE    35     -14.683  14.059   7.205  0.50213.39       1SG 284
     ATOM    284  CG  PHE    35     -15.624  14.612   8.170  0.50213.39       1SG 285
     ATOM    285  CD1 PHE    35     -15.385  15.592   9.112  0.50213.39       1SG 286
     ATOM    286  CD2 PHE    35     -16.780  13.907   8.118  0.50213.39       1SG 287
     ATOM    287  CE1 PHE    35     -16.409  15.856   9.992  0.50213.39       1SG 288
     ATOM    288  CE2 PHE    35     -17.793  14.160   8.976  0.50213.39       1SG 289
     ATOM    289  CZ  PHE    35     -17.589  15.139   9.915  0.50213.39       1SG 290
     ATOM    290  C   PHE    35     -13.342  13.381   5.900  1.00213.39       1SG 291
     ATOM    291  O   PHE    35     -12.488  13.238   5.039  1.00213.39       1SG 292
*cap ATOM    292  NT  PHE    35     -13.553  12.516   6.942  1.00213.39       1SG 293
     END

"""

    # Copy the original PDB to a temp file, to preserve a record off it
    tmpPdbFile = outPdbFile + '.before'
    fin = open(outPdbFile,'r')
    fout = open(tmpPdbFile, 'w')
    fout.write( fin.read() )
    fin.close()
    fout.close()

    # Read in the atoms from the PDB
    atoms = readAtomsFromPDB( outPdbFile )

    # Find the atoms in the termini caps
    nter_atoms = []
    cter_atoms = []
    i=0
    while i < len(atoms):
        if atoms[i]['name'].strip() == 'CY':
            atoms[i]['name'] = ' C  '
            nter_atoms.insert(0, atoms.pop(i))
        elif atoms[i]['name'].strip() == 'CAY':
            atoms[i]['name'] = ' CA '
            nter_atoms.insert(0, atoms.pop(i) )
        elif atoms[i]['name'].strip() == 'OY':
            atoms[i]['name'] = ' O  '
            nter_atoms.insert(0, atoms.pop(i) )
        elif atoms[i]['name'].strip() == 'NT':
            atoms[i]['name'] = ' N  '
            cter_atoms.insert(0, atoms.pop(i) )
        else:
            i += 1

    # print 'nter_atoms', nter_atoms, 'cter_atoms', cter_atoms

    # Rename, shift numbering, and tack back onto the atoms list
    for atom in nter_atoms:
        atom['resName'] = 'ACE '
        atom['resSeq'] = atom['resSeq'] - 1
        atoms.insert(0, atom)
    for atom in cter_atoms:
        atom['resName'] = 'NH2 '
        atom['resSeq'] = atom['resSeq'] + 1
        atoms.append( atom )

    # write he atoms to PDB -- this will take care of renumbering too
    writeAtomsToPDB(outPdbFile, atoms, renumber = True)

<|endoftext|>


<commit_before>from hwt.hdlObjects.constants import INTF_DIRECTION
from hwt.synthesizer.param import evalParam

def autoAddAgents(unit):
    """
    Walk all interfaces on unit and instantiate actor for every interface.

    :return: all monitor/driver functions which should be added to simulation as processes
    """
    proc = []
    for intf in unit._interfaces:
        if not intf._isExtern:
            continue

        try:
            agentCls = intf._getSimAgent()
        except NotImplementedError:
            raise NotImplementedError(("Interface %s\n" +
                                       "has not any simulation agent class assigned") % (str(intf)))

        if intf._isInterfaceArray():
            agentCnt = evalParam(intf._multipliedBy).val
            agent = []
            for i in range(agentCnt):
                _intf = intf[i]
                a = agentCls(_intf)
                agent.append(a)
                _intf._ag = a
        else:
            agent = agentCls(intf)
            intf._ag = agent

        if intf._multipliedBy is None:
            agent = [agent, ]

        if intf._direction == INTF_DIRECTION.MASTER:
            agProcs = list(map(lambda a: a.getMonitors(), agent))
        elif intf._direction == INTF_DIRECTION.SLAVE:
            agProcs = list(map(lambda a: a.getDrivers(), agent))
        else:
            raise NotImplementedError("intf._direction %s for %r" % (str(intf._direction), intf))

        for p in agProcs:
            proc.extend(p)

    return proc

def valuesToInts(values):
    """
    Iterable of values to ints (nonvalid = None)
    """
    res = []
    for d in values:
        res.append(valToInt(d))
    return res

def valToInt(v):
    if v.vldMask == v._dtype.all_mask():
        return v.val
    else:
        return None

def agInts(interface):
    """
    Convert all values which has agent collected in time >=0 to integer array.
    Invalid value will be None.
    """
    return valuesToInts(interface._ag.data)
<commit_msg>grammar fix<commit_after>from hwt.hdlObjects.constants import INTF_DIRECTION
from hwt.synthesizer.param import evalParam

def autoAddAgents(unit):
    """
    Walk all interfaces on unit and instantiate agent for every interface.

    :return: all monitor/driver functions which should be added to simulation as processes
    """
    proc = []
    for intf in unit._interfaces:
        if not intf._isExtern:
            continue

        try:
            agentCls = intf._getSimAgent()
        except NotImplementedError:
            raise NotImplementedError(("Interface %s\n" +
                                       "has not any simulation agent class assigned") % (str(intf)))

        if intf._isInterfaceArray():
            agentCnt = evalParam(intf._multipliedBy).val
            agent = []
            for i in range(agentCnt):
                _intf = intf[i]
                a = agentCls(_intf)
                agent.append(a)
                _intf._ag = a
        else:
            agent = agentCls(intf)
            intf._ag = agent

        if intf._multipliedBy is None:
            agent = [agent, ]

        if intf._direction == INTF_DIRECTION.MASTER:
            agProcs = list(map(lambda a: a.getMonitors(), agent))
        elif intf._direction == INTF_DIRECTION.SLAVE:
            agProcs = list(map(lambda a: a.getDrivers(), agent))
        else:
            raise NotImplementedError("intf._direction %s for %r" % (str(intf._direction), intf))

        for p in agProcs:
            proc.extend(p)

    return proc

def valuesToInts(values):
    """
    Iterable of values to ints (nonvalid = None)
    """
    res = []
    for d in values:
        res.append(valToInt(d))
    return res

def valToInt(v):
    if v.vldMask == v._dtype.all_mask():
        return v.val
    else:
        return None

def agInts(interface):
    """
    Convert all values which has agent collected in time >=0 to integer array.
    Invalid value will be None.
    """
    return valuesToInts(interface._ag.data)
<|endoftext|>


<commit_before>def _pad_binary(bin_str, req_len=8):
    return max(0, req_len - len(bin_str)) * '0' + bin_str

def _hex_to_bin_str(hex_string):
    unpadded_bin_string_list = map(bin, hex_string)
    padded_bin_string_list = map(_pad_binary, unpadded_bin_string_list)
    bitwise_message = "".join(padded_bin_string_list)
    return bitwise_message

class HuffmanDecoder(object):
    class _Node(object):
        def __init__(self, data):
            self.data = data
            self.mapping = {}

    def __init__(self, huffman_code_list, huffman_code_list_lengths):
        self.root = self._Node(None)
        for index, (huffman_code, code_length) in enumerate(zip(huffman_code_list, huffman_code_list_lengths)):
            self._insert(huffman_code, code_length, chr(index))

    def _insert(self, hex_number, hex_length, letter):
        hex_number = _pad_binary(bin(hex_number)[2:], hex_length)
        cur_node = self.root
        for digit in hex_number:
            if digit not in cur_node.mapping:
                cur_node.mapping[digit] = self._Node(None)
            cur_node = cur_node.mapping[digit]
        cur_node.data = letter

    def decode(self, encoded_string):
        number = _hex_to_bin_str(encoded_string)
        cur_node = self.root
        decoded_message = []
        for digit in number:
            if digit not in cur_node.mapping:
                break
            cur_node = cur_node.mapping[digit]
            if cur_node.data is not None:
                decoded_message.append(cur_node.data)
                cur_node = self.root
        return "".join(decoded_message)
<commit_msg>Move the removal of the binary prefix.<commit_after>def _pad_binary(bin_str, req_len=8):
    bin_str = bin_str[2:]  # Strip the 0b prefix
    return max(0, req_len - len(bin_str)) * '0' + bin_str

def _hex_to_bin_str(hex_string):
    unpadded_bin_string_list = map(bin, hex_string)
    padded_bin_string_list = map(_pad_binary, unpadded_bin_string_list)
    bitwise_message = "".join(padded_bin_string_list)
    return bitwise_message

class HuffmanDecoder(object):
    class _Node(object):
        def __init__(self, data):
            self.data = data
            self.mapping = {}

    def __init__(self, huffman_code_list, huffman_code_list_lengths):
        self.root = self._Node(None)
        for index, (huffman_code, code_length) in enumerate(zip(huffman_code_list, huffman_code_list_lengths)):
            self._insert(huffman_code, code_length, chr(index))

    def _insert(self, hex_number, hex_length, letter):
        hex_number = _pad_binary(bin(hex_number), hex_length)
        cur_node = self.root
        for digit in hex_number:
            if digit not in cur_node.mapping:
                cur_node.mapping[digit] = self._Node(None)
            cur_node = cur_node.mapping[digit]
        cur_node.data = letter

    def decode(self, encoded_string):
        number = _hex_to_bin_str(encoded_string)
        cur_node = self.root
        decoded_message = []
        for digit in number:
            if digit not in cur_node.mapping:
                break
            cur_node = cur_node.mapping[digit]
            if cur_node.data is not None:
                decoded_message.append(cur_node.data)
                cur_node = self.root
        return "".join(decoded_message)
<|endoftext|>


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

"""

These are System tests

They will only work if you have all the moving parts available

"""

import unittest2
import guestfs
import os
import time
import subprocess
import threading
import sqlite3
import StringIO

from . import core
from . import sql
from . import monitor

class HyperkitTestError(Exception):
    pass

class TestRunFailed(Exception):
    pass

class TestRun(object):

    timeout = 120

    def __init__(self, directory, hypervisor, distro, release, arch):
        self.directory = directory
        self.hypervisor = hypervisor
        self.distro = distro
        self.release = release
        self.arch = arch
        self.results = {}
        self.executing = None
        self.guest = None
        self.timer = None

    def __enter__(self):
        self.create()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.stop_timer()
        self.terminate()
        self.cleanup()
        if self.guest is not None:
            self.guest.shutdown()

    @property
    def username(self):
        return "XXusernameXX"

    @property
    def password(self):
        return "XXpasswordXX"

    @property
    def instance(self):
        return "-".join([self.hypervisor, self.distro, self.release, self.arch])

    @property
    def name(self):
        return "system_test-" + self.instance

    @property
    def public_key(self):
        keyname = os.path.join(self.directory, self.name +".ssh_key")
        if not os.path.exists(keyname):
            print >> open(keyname, "w"), "FAKE PUBLIC KEY"
        return keyname

    @property
    def path(self):
        return os.path.join(self.directory, self.name)

    def create(self):
        if os.path.exists(self.path):
            raise KeyError("Path %r already exists" % self.path)
        args = [
            "--username", self.username,
            "--password", self.password,
            #"--public-key", self.public_key,
            self.name, self.distro, self.release, self.arch
        ]
        r = self.hyperkit("create", args)
        self.results['create'] = r
        if r['code'] != 0:
            raise TestRunFailed

    def cleanup(self):
        self.hyperkit("cleanup", ["--yes", self.name])

    def terminate(self):
        if self.executing is not None and self.executing.poll() is None:
            print "Test taking too long, terminating running process"
            try:
                self.executing.kill()
            except:
                # ignore race condition
                pass

    def start_timer(self):
        self.start_time = time.time()
        self.timer = threading.Timer(self.timeout, self.terminate)
        self.timer.start()

    def stop_timer(self):
        if self.timer is not None:
            self.timer.cancel()

    def run_tests(self):
        self.run_stage("start")
        self.run_stage("wait")
        self.run_stage("stop")
        self.analyse()
        return self.results

    def run_stage(self, stage):
        self.start_timer()
        r = self.hyperkit(stage, [self.name])
        self.stop_timer()
        self.results[stage] = r

    def hyperkit(self, command, args):
        command = ["hyperkit",
                   "--debug",
                   "--directory", self.directory,
                   "--hypervisor", self.hypervisor,
                   command,
                   ] + args
        print "Executing", " ".join(command)
        t = time.time()
        self.executing = subprocess.Popen(command,
                                          stdout=subprocess.PIPE,
                                          stderr=subprocess.PIPE)
        stdout, stderr = self.executing.communicate()
        return_code = self.executing.returncode
        self.executing = None
        print "Command completed in %0.2fs" % (time.time() - t)
        return dict(stdout=stdout, stderr=stderr, code=return_code)

    def mount_guest(self):
        disk_name = self.name + "_disk1.vdi"  # conditional on hypervisor
        disk_path = os.path.join(self.directory, self.name, disk_name)
        self.guest = guestfs.GuestFS(python_return_dict=True)
        self.guest.add_drive_opts(disk_path, readonly=1)
        self.guest.launch()
        self.guest.mount_ro("/dev/sda1", "/")

    def analyse(self):
        """ Analyze a disk image """
        self.mount_guest()
        loader = unittest2.TestLoader()
        # yes this is horrific. i will fix it at some point.
        core.guest = self.guest
        suite = loader.loadTestsFromModule(core)
        results = StringIO.StringIO()
        unittest2.TextTestRunner(stream=results, verbosity=2).run(suite)
        self.results["analysis"] = results.getvalue()

class SystemTestManager:

    db_name = "hyperkit.sqlite"
    hypervisors = ["vbox", "vmware"]

    distros = [{
        "name": "ubuntu",
        "releases": ["14.04", "12.04"],
        "architectures": ["amd64", "i386"],
    }, {
        "name": "fedora",
        "releases": ["20", "19"],
        "architectures": ["x86", "x86_64"],
    }]

    def __init__(self, directory):
        directory = os.path.realpath(directory)
        dbpath = os.path.join(directory, self.db_name)
        if not os.path.exists(dbpath):
            print "Database %s does not exist - have you specified the test directory?" % dbpath
            raise SystemExit
        self.directory = directory
        self.db = sqlite3.connect(dbpath)
        self.executing = None

    def set_hypervisors(self, hypervisors):
        c = self.cursor
        c.execute("delete from hypervisor")
        for h in hypervisors:
            c.execute(sql.hypervisor_insert, [h])
        self.db.commit()

    def set_releases(self, releases):
        for r in releases:
            if not self.check_release(*r):
                print "Bad release spec:", r
                raise SystemExit
        c = self.cursor
        c.execute("delete from candidate")
        for d, r, a in releases:
            c.execute(sql.candidate_insert, [d, r, a])
        self.db.commit()

    @property
    def cursor(self):
        return self.db.cursor()

    @classmethod
    def create(self, directory):
        if not os.path.exists(directory):
            os.mkdir(directory)
        dbpath = os.path.join(directory, self.db_name)
        print "Creating test database", dbpath
        db = sqlite3.connect(dbpath)
        cur = db.cursor()
        cur.executescript(sql.create)
        m = self(directory)
        m.set_hypervisors(self.hypervisors)
        c = m.cursor
        for d in m.distros:
            distro = d['name']
            c.execute(sql.distro_insert, [distro])
            for r in d['releases']:
                c.execute(sql.release_insert, (r, distro))
            for a in d['architectures']:
                c.execute(sql.architecture_insert, (a, distro))
                for r in d['releases']:
                    c.execute(sql.candidate_insert, (distro, r, a))
        m.db.commit()

    def check_release(self, distro, release=None, architecture=None):
        for e in self.distros:
            if e['name'] == distro:
                if release not in e['releases']:
                    print "Release", release, "not recognised for distro", distro
                    return False
                if architecture not in e['architectures']:
                    print "Architecture", architecture, "not recognised for distro", distro
                    return False
                return True
        return False

    def should_test(self, hypervisor, distro, release, architecture):
        if self.hypervisors and hypervisor not in self.hypervisors:
            return False
        if self.releases and (distro, release, architecture) not in self.releases:
            return False
        return True

    def test_kernel_permissions(self):
        p = subprocess.Popen(["uname", "-r"], stdout=subprocess.PIPE)
        release = p.stdout.read().strip()
        kernel = "/boot/vmlinuz-%s" % release
        if not os.access(kernel, os.R_OK):
            raise HyperkitTestError("Cannot read kernel %s. You need to chmod +r this file." % kernel)

    def exec_test(self, hypervisor, distro, release, arch):
        self.test_kernel_permissions()
        with TestRun(self.directory, hypervisor, distro, release, arch) as r:
            results = r.run_tests()
            return results
<commit_msg>reimplemented the analysis as a context manager too<commit_after>#! /usr/bin/env python

"""

These are System tests

They will only work if you have all the moving parts available

"""

import unittest2
import guestfs
import os
import time
import subprocess
import threading
import sqlite3
import StringIO

from . import core
from . import sql
from . import monitor

class HyperkitTestError(Exception):
    pass

class TestRunFailed(Exception):
    pass

class GuestAnalysis(object):

    def __init__(self, test_run):
        self.test_run = test_run

    def __enter__(self):
        disk_name = self.test_run.name + "_disk1.vdi"  # conditional on hypervisor
        disk_path = os.path.join(self.test_run.directory, self.test_run.name, disk_name)
        self.guest = guestfs.GuestFS(python_return_dict=True)
        self.guest.add_drive_opts(disk_path, readonly=1)
        self.guest.launch()
        self.guest.mount_ro("/dev/sda1", "/")
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.guest.shutdown()

    def analyse(self):
        """ Analyze a disk image """
        loader = unittest2.TestLoader()
        # yes this is horrific. i will fix it at some point.
        core.guest = self.guest
        suite = loader.loadTestsFromModule(core)
        results = StringIO.StringIO()
        unittest2.TextTestRunner(stream=results, verbosity=2).run(suite)
        return results.getvalue()

class TestRun(object):

    timeout = 120

    def __init__(self, directory, hypervisor, distro, release, arch):
        self.directory = directory
        self.hypervisor = hypervisor
        self.distro = distro
        self.release = release
        self.arch = arch
        self.results = {}
        self.executing = None
        self.guest = None
        self.timer = None

    def __enter__(self):
        self.create()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.stop_timer()
        self.terminate()
        self.cleanup()

    @property
    def username(self):
        return "XXusernameXX"

    @property
    def password(self):
        return "XXpasswordXX"

    @property
    def instance(self):
        return "-".join([self.hypervisor, self.distro, self.release, self.arch])

    @property
    def name(self):
        return "system_test-" + self.instance

    @property
    def public_key(self):
        keyname = os.path.join(self.directory, self.name +".ssh_key")
        if not os.path.exists(keyname):
            print >> open(keyname, "w"), "FAKE PUBLIC KEY"
        return keyname

    @property
    def path(self):
        return os.path.join(self.directory, self.name)

    def create(self):
        if os.path.exists(self.path):
            raise KeyError("Path %r already exists" % self.path)
        args = [
            "--username", self.username,
            "--password", self.password,
            #"--public-key", self.public_key,
            self.name, self.distro, self.release, self.arch
        ]
        r = self.hyperkit("create", args)
        self.results['create'] = r
        if r['code'] != 0:
            raise TestRunFailed

    def cleanup(self):
        self.hyperkit("cleanup", ["--yes", self.name])

    def terminate(self):
        if self.executing is not None and self.executing.poll() is None:
            print "Test taking too long, terminating running process"
            try:
                self.executing.kill()
            except:
                # ignore race condition
                pass

    def start_timer(self):
        self.start_time = time.time()
        self.timer = threading.Timer(self.timeout, self.terminate)
        self.timer.start()

    def stop_timer(self):
        if self.timer is not None:
            self.timer.cancel()

    def run_tests(self):
        self.run_stage("start")
        self.run_stage("wait")
        self.run_stage("stop")
        self.analyse()
        return self.results

    def run_stage(self, stage):
        self.start_timer()
        r = self.hyperkit(stage, [self.name])
        self.stop_timer()
        self.results[stage] = r

    def hyperkit(self, command, args):
        command = ["hyperkit",
                   "--debug",
                   "--directory", self.directory,
                   "--hypervisor", self.hypervisor,
                   command,
                   ] + args
        print "Executing", " ".join(command)
        t = time.time()
        self.executing = subprocess.Popen(command,
                                          stdout=subprocess.PIPE,
                                          stderr=subprocess.PIPE)
        stdout, stderr = self.executing.communicate()
        return_code = self.executing.returncode
        self.executing = None
        print "Command completed in %0.2fs" % (time.time() - t)
        return dict(stdout=stdout, stderr=stderr, code=return_code)

    def analyse(self):
        with GuestAnalysis(self) as analyser:
            self.results['analysis'] = analyser.analyse()

class SystemTestManager:

    db_name = "hyperkit.sqlite"
    hypervisors = ["vbox", "vmware"]

    distros = [{
        "name": "ubuntu",
        "releases": ["14.04", "12.04"],
        "architectures": ["amd64", "i386"],
    }, {
        "name": "fedora",
        "releases": ["20", "19"],
        "architectures": ["x86", "x86_64"],
    }]

    def __init__(self, directory):
        directory = os.path.realpath(directory)
        dbpath = os.path.join(directory, self.db_name)
        if not os.path.exists(dbpath):
            print "Database %s does not exist - have you specified the test directory?" % dbpath
            raise SystemExit
        self.directory = directory
        self.db = sqlite3.connect(dbpath)
        self.executing = None

    def set_hypervisors(self, hypervisors):
        c = self.cursor
        c.execute("delete from hypervisor")
        for h in hypervisors:
            c.execute(sql.hypervisor_insert, [h])
        self.db.commit()

    def set_releases(self, releases):
        for r in releases:
            if not self.check_release(*r):
                print "Bad release spec:", r
                raise SystemExit
        c = self.cursor
        c.execute("delete from candidate")
        for d, r, a in releases:
            c.execute(sql.candidate_insert, [d, r, a])
        self.db.commit()

    @property
    def cursor(self):
        return self.db.cursor()

    @classmethod
    def create(self, directory):
        if not os.path.exists(directory):
            os.mkdir(directory)
        dbpath = os.path.join(directory, self.db_name)
        print "Creating test database", dbpath
        db = sqlite3.connect(dbpath)
        cur = db.cursor()
        cur.executescript(sql.create)
        m = self(directory)
        m.set_hypervisors(self.hypervisors)
        c = m.cursor
        for d in m.distros:
            distro = d['name']
            c.execute(sql.distro_insert, [distro])
            for r in d['releases']:
                c.execute(sql.release_insert, (r, distro))
            for a in d['architectures']:
                c.execute(sql.architecture_insert, (a, distro))
                for r in d['releases']:
                    c.execute(sql.candidate_insert, (distro, r, a))
        m.db.commit()

    def check_release(self, distro, release=None, architecture=None):
        for e in self.distros:
            if e['name'] == distro:
                if release not in e['releases']:
                    print "Release", release, "not recognised for distro", distro
                    return False
                if architecture not in e['architectures']:
                    print "Architecture", architecture, "not recognised for distro", distro
                    return False
                return True
        return False

    def should_test(self, hypervisor, distro, release, architecture):
        if self.hypervisors and hypervisor not in self.hypervisors:
            return False
        if self.releases and (distro, release, architecture) not in self.releases:
            return False
        return True

    def test_kernel_permissions(self):
        p = subprocess.Popen(["uname", "-r"], stdout=subprocess.PIPE)
        release = p.stdout.read().strip()
        kernel = "/boot/vmlinuz-%s" % release
        if not os.access(kernel, os.R_OK):
            raise HyperkitTestError("Cannot read kernel %s. You need to chmod +r this file." % kernel)

    def exec_test(self, hypervisor, distro, release, arch):
        self.test_kernel_permissions()
        with TestRun(self.directory, hypervisor, distro, release, arch) as r:
            results = r.run_tests()
            return results
<|endoftext|>


<commit_before># -*- coding: iso-8859-15 -*-
# =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2016 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import inspect
import logging

from django.db import connection
from django.db.models import Max, Min, Count
from django.conf import settings

from pycsw.core import util
from hypermap.aggregator.models import Catalog, Layer, Service, Endpoint, SpatialReferenceSystem
from hypermap.aggregator.utils import create_layer_from_metadata_xml
from hypermap.aggregator.tasks import index_layer

LOGGER = logging.getLogger(__name__)

HYPERMAP_SERVICE_TYPES = {
    # 'HHypermap enum': 'CSW enum'
    'http://www.opengis.net/cat/csw/2.0.2': 'OGC:CSW',
    'http://www.opengis.net/wms': 'OGC:WMS',
    'http://www.opengis.net/wmts/1.0': 'OGC:WMTS',
    'https://wiki.osgeo.org/wiki/TMS': 'OSGeo:TMS',
    'urn:x-esri:serviceType:ArcGIS:MapServer': 'ESRI:ArcGIS:MapServer',
    'urn:x-esri:serviceType:ArcGIS:ImageServer': 'ESRI:ArcGIS:ImageServer'
}

def get_service(raw_xml):
    """
    Set a service object based on the XML metadata
       <dct:references scheme="OGC:WMS">http://ngamaps.geointapps.org/arcgis
       /services/RIO/Rio_Foundation_Transportation/MapServer/WMSServer
       </dct:references>
    :param instance:
    :return: Layer
    """
    from pycsw.core.etree import etree

    parsed = etree.fromstring(raw_xml, etree.XMLParser(resolve_entities=False))

    # <dc:format>OGC:WMS</dc:format>
    source_tag = parsed.find("{http://purl.org/dc/elements/1.1/}source")
    # <dc:source>
    #    http://ngamaps.geointapps.org/arcgis/services/RIO/Rio_Foundation_Transportation/MapServer/WMSServer
    # </dc:source>
    format_tag = parsed.find("{http://purl.org/dc/elements/1.1/}format")

    service_url = None
    service_type = None

    if hasattr(source_tag, 'text'):
        service_url = source_tag.text

    if hasattr(format_tag, 'text'):
        service_type = format_tag.text

    if hasattr(format_tag, 'text'):
        service_type = format_tag.text

    service, created = Service.objects.get_or_create(url=service_url,
                                                     is_monitored=False,
                                                     type=service_type)
    # TODO: dont hardcode SRS, get them from the parsed XML.
    srs, created = SpatialReferenceSystem.objects.get_or_create(code="EPSG:4326")
    service.srs.add(srs)

    return service

class HHypermapRepository(object):
    ''' Class to interact with underlying repository '''
    def __init__(self, context, repo_filter=None):
        ''' Initialize repository '''

        self.context = context
        self.filter = repo_filter
        self.fts = False
        self.label = 'HHypermap'
        self.local_ingest = True

        self.dbtype = settings.DATABASES['default']['ENGINE'].split('.')[-1]

        # HHypermap PostgreSQL installs are PostGIS enabled
        if self.dbtype == 'postgis':
            self.dbtype = 'postgresql+postgis+wkt'

        if self.dbtype in ['sqlite', 'sqlite3']:  # load SQLite query bindings
            connection.connection.create_function('query_spatial', 4, util.query_spatial)
            connection.connection.create_function('get_anytext', 1, util.get_anytext)
            connection.connection.create_function('get_geometry_area', 1, util.get_geometry_area)

        # generate core queryables db and obj bindings
        self.queryables = {}

        for tname in self.context.model['typenames']:
            for qname in self.context.model['typenames'][tname]['queryables']:
                self.queryables[qname] = {}
                items = self.context.model['typenames'][tname]['queryables'][qname].items()

                for qkey, qvalue in items:
                    self.queryables[qname][qkey] = qvalue

        # flatten all queryables
        # TODO smarter way of doing this
        self.queryables['_all'] = {}
        for qbl in self.queryables:
            self.queryables['_all'].update(self.queryables[qbl])
        self.queryables['_all'].update(self.context.md_core_model['mappings'])

        if 'Harvest' in self.context.model['operations'] and 'Transaction' in self.context.model['operations']:
            self.context.model['operations']['Harvest']['parameters']['ResourceType']['values'] = HYPERMAP_SERVICE_TYPES.keys()  # noqa
            self.context.model['operations']['Transaction']['parameters']['TransactionSchemas']['values'] = HYPERMAP_SERVICE_TYPES.keys()  # noqa

    def dataset(self):
        ''' Stub to mock a pycsw dataset object for Transactions'''
        return type('Service', (object,), {})

    def query_ids(self, ids):
        ''' Query by list of identifiers '''
        return self._get_repo_filter(Layer.objects).filter(id__in=ids).all()

    def query_domain(self, domain, typenames, domainquerytype='list', count=False):
        ''' Query by property domain values '''

        objects = self._get_repo_filter(Layer.objects)

        if domainquerytype == 'range':
            return [tuple(objects.aggregate(Min(domain), Max(domain)).values())]
        else:
            if count:
                return [(d[domain], d['%s__count' % domain])
                        for d in objects.values(domain).annotate(Count(domain))]
            else:
                return objects.values_list(domain).distinct()

    def query_insert(self, direction='max'):
        ''' Query to get latest (default) or earliest update to repository '''
        if direction == 'min':
            return Layer.objects.aggregate(
                Min('last_updated'))['last_updated__min'].strftime('%Y-%m-%dT%H:%M:%SZ')
        return self._get_repo_filter(Layer.objects).aggregate(
            Max('last_updated'))['last_updated__max'].strftime('%Y-%m-%dT%H:%M:%SZ')

    def query_source(self, source):
        ''' Query by source '''
        return self._get_repo_filter(Layer.objects).filter(url=source)

    def query(self, constraint, sortby=None, typenames=None, maxrecords=10, startposition=0):
        ''' Query records from underlying repository '''

        # run the raw query and get total
        if 'where' in constraint:  # GetRecords with constraint
            query = self._get_repo_filter(Layer.objects).extra(where=[constraint['where']], params=constraint['values'])

        else:  # GetRecords sans constraint
            query = self._get_repo_filter(Layer.objects)

        total = query.count()

        # apply sorting, limit and offset
        if sortby is not None:
            if 'spatial' in sortby and sortby['spatial']:  # spatial sort
                desc = False
                if sortby['order'] == 'DESC':
                    desc = True
                query = query.all()
                return [str(total),
                        sorted(query,
                               key=lambda x: float(util.get_geometry_area(getattr(x, sortby['propertyname']))),
                               reverse=desc,
                               )[startposition:startposition+int(maxrecords)]]
            else:
                if sortby['order'] == 'DESC':
                    pname = '-%s' % sortby['propertyname']
                else:
                    pname = sortby['propertyname']
                return [str(total),
                        query.order_by(pname)[startposition:startposition+int(maxrecords)]]
        else:  # no sort
            return [str(total), query.all()[startposition:startposition+int(maxrecords)]]

    def insert(self, resourcetype, source, insert_date=None):
        ''' Insert a record into the repository '''

        caller = inspect.stack()[1][3]

        if caller == 'transaction':  # insert of Layer
            hhclass = 'Layer'
            source = resourcetype
            resourcetype = resourcetype.csw_schema
        else:  # insert of service
            hhclass = 'Service'
            if resourcetype not in HYPERMAP_SERVICE_TYPES.keys():
                raise RuntimeError('Unsupported Service Type')

        return self._insert_or_update(resourcetype, source, mode='insert', hhclass=hhclass)

    def _insert_or_update(self, resourcetype, source, mode='insert', hhclass='Service'):
        ''' Insert or update a record in the repository '''

        keywords = []

        if self.filter is not None:
            catalog = Catalog.objects.get(id=int(self.filter.split()[-1]))

        try:
            if hhclass == 'Layer':
                # TODO: better way of figuring out duplicates
                match = Layer.objects.filter(name=source.name,
                                             title=source.title,
                                             abstract=source.abstract,
                                             is_monitored=False)
                matches = match.all()
                if matches:
                    if mode == 'insert':
                        raise RuntimeError('HHypermap error: Layer %d \'%s\' already exists' % (
                                           matches[0].id, source.title))
                    elif mode == 'update':
                        match.update(
                            name=source.name,
                            title=source.title,
                            abstract=source.abstract,
                            is_monitored=False,
                            xml=source.xml,
                            wkt_geometry=source.wkt_geometry,
                            anytext=util.get_anytext([source.title, source.abstract, source.keywords_csv])
                        )

                service = get_service(source.xml)
                res, keywords = create_layer_from_metadata_xml(resourcetype, source.xml,
                                                               monitor=False, service=service,
                                                               catalog=catalog)

                res.save()

                LOGGER.debug('Indexing layer with id %s on search engine' % res.id)
                index_layer(res, use_cache=True)

            else:
                if resourcetype == 'http://www.opengis.net/cat/csw/2.0.2':
                    res = Endpoint(url=source, catalog=catalog)
                else:
                    res = Service(type=HYPERMAP_SERVICE_TYPES[resourcetype], url=source, catalog=catalog)

                res.save()

            if keywords:
                for kw in keywords:
                    res.keywords.add(kw)
        except Exception as err:
            raise RuntimeError('HHypermap error: %s' % err)

        # return a list of ids that were inserted or updated
        ids = []

        if hhclass == 'Layer':
            ids.append({'identifier': res.uuid, 'title': res.title})
        else:
            if resourcetype == 'http://www.opengis.net/cat/csw/2.0.2':
                for res in Endpoint.objects.filter(url=source).all():
                    ids.append({'identifier': res.uuid, 'title': res.url})
            else:
                for res in Service.objects.filter(url=source).all():
                    ids.append({'identifier': res.uuid, 'title': res.title})

        return ids

    def delete(self, constraint):
        ''' Delete a record from the repository '''

        results = self._get_repo_filter(Service.objects).extra(where=[constraint['where']],
                                                               params=constraint['values']).all()
        deleted = len(results)
        results.delete()
        return deleted

    def _get_repo_filter(self, query):
        ''' Apply repository wide side filter / mask query '''
        if self.filter is not None:
            return query.extra(where=[self.filter])
        return query
<commit_msg>We need to switch from id to uuid in the pycsw plugin<commit_after># -*- coding: iso-8859-15 -*-
# =================================================================
#
# Authors: Tom Kralidis <tomkralidis@gmail.com>
#
# Copyright (c) 2016 Tom Kralidis
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
# =================================================================

import inspect
import logging

from django.db import connection
from django.db.models import Max, Min, Count
from django.conf import settings

from pycsw.core import util
from hypermap.aggregator.models import Catalog, Layer, Service, Endpoint, SpatialReferenceSystem
from hypermap.aggregator.utils import create_layer_from_metadata_xml
from hypermap.aggregator.tasks import index_layer

LOGGER = logging.getLogger(__name__)

HYPERMAP_SERVICE_TYPES = {
    # 'HHypermap enum': 'CSW enum'
    'http://www.opengis.net/cat/csw/2.0.2': 'OGC:CSW',
    'http://www.opengis.net/wms': 'OGC:WMS',
    'http://www.opengis.net/wmts/1.0': 'OGC:WMTS',
    'https://wiki.osgeo.org/wiki/TMS': 'OSGeo:TMS',
    'urn:x-esri:serviceType:ArcGIS:MapServer': 'ESRI:ArcGIS:MapServer',
    'urn:x-esri:serviceType:ArcGIS:ImageServer': 'ESRI:ArcGIS:ImageServer'
}

def get_service(raw_xml):
    """
    Set a service object based on the XML metadata
       <dct:references scheme="OGC:WMS">http://ngamaps.geointapps.org/arcgis
       /services/RIO/Rio_Foundation_Transportation/MapServer/WMSServer
       </dct:references>
    :param instance:
    :return: Layer
    """
    from pycsw.core.etree import etree

    parsed = etree.fromstring(raw_xml, etree.XMLParser(resolve_entities=False))

    # <dc:format>OGC:WMS</dc:format>
    source_tag = parsed.find("{http://purl.org/dc/elements/1.1/}source")
    # <dc:source>
    #    http://ngamaps.geointapps.org/arcgis/services/RIO/Rio_Foundation_Transportation/MapServer/WMSServer
    # </dc:source>
    format_tag = parsed.find("{http://purl.org/dc/elements/1.1/}format")

    service_url = None
    service_type = None

    if hasattr(source_tag, 'text'):
        service_url = source_tag.text

    if hasattr(format_tag, 'text'):
        service_type = format_tag.text

    if hasattr(format_tag, 'text'):
        service_type = format_tag.text

    service, created = Service.objects.get_or_create(url=service_url,
                                                     is_monitored=False,
                                                     type=service_type)
    # TODO: dont hardcode SRS, get them from the parsed XML.
    srs, created = SpatialReferenceSystem.objects.get_or_create(code="EPSG:4326")
    service.srs.add(srs)

    return service

class HHypermapRepository(object):
    ''' Class to interact with underlying repository '''
    def __init__(self, context, repo_filter=None):
        ''' Initialize repository '''

        self.context = context
        self.filter = repo_filter
        self.fts = False
        self.label = 'HHypermap'
        self.local_ingest = True

        self.dbtype = settings.DATABASES['default']['ENGINE'].split('.')[-1]

        # HHypermap PostgreSQL installs are PostGIS enabled
        if self.dbtype == 'postgis':
            self.dbtype = 'postgresql+postgis+wkt'

        if self.dbtype in ['sqlite', 'sqlite3']:  # load SQLite query bindings
            connection.connection.create_function('query_spatial', 4, util.query_spatial)
            connection.connection.create_function('get_anytext', 1, util.get_anytext)
            connection.connection.create_function('get_geometry_area', 1, util.get_geometry_area)

        # generate core queryables db and obj bindings
        self.queryables = {}

        for tname in self.context.model['typenames']:
            for qname in self.context.model['typenames'][tname]['queryables']:
                self.queryables[qname] = {}
                items = self.context.model['typenames'][tname]['queryables'][qname].items()

                for qkey, qvalue in items:
                    self.queryables[qname][qkey] = qvalue

        # flatten all queryables
        # TODO smarter way of doing this
        self.queryables['_all'] = {}
        for qbl in self.queryables:
            self.queryables['_all'].update(self.queryables[qbl])
        self.queryables['_all'].update(self.context.md_core_model['mappings'])

        if 'Harvest' in self.context.model['operations'] and 'Transaction' in self.context.model['operations']:
            self.context.model['operations']['Harvest']['parameters']['ResourceType']['values'] = HYPERMAP_SERVICE_TYPES.keys()  # noqa
            self.context.model['operations']['Transaction']['parameters']['TransactionSchemas']['values'] = HYPERMAP_SERVICE_TYPES.keys()  # noqa

    def dataset(self):
        ''' Stub to mock a pycsw dataset object for Transactions'''
        return type('Service', (object,), {})

    def query_ids(self, ids):
        ''' Query by list of identifiers '''
        return self._get_repo_filter(Layer.objects).filter(uuid__in=ids).all()

    def query_domain(self, domain, typenames, domainquerytype='list', count=False):
        ''' Query by property domain values '''

        objects = self._get_repo_filter(Layer.objects)

        if domainquerytype == 'range':
            return [tuple(objects.aggregate(Min(domain), Max(domain)).values())]
        else:
            if count:
                return [(d[domain], d['%s__count' % domain])
                        for d in objects.values(domain).annotate(Count(domain))]
            else:
                return objects.values_list(domain).distinct()

    def query_insert(self, direction='max'):
        ''' Query to get latest (default) or earliest update to repository '''
        if direction == 'min':
            return Layer.objects.aggregate(
                Min('last_updated'))['last_updated__min'].strftime('%Y-%m-%dT%H:%M:%SZ')
        return self._get_repo_filter(Layer.objects).aggregate(
            Max('last_updated'))['last_updated__max'].strftime('%Y-%m-%dT%H:%M:%SZ')

    def query_source(self, source):
        ''' Query by source '''
        return self._get_repo_filter(Layer.objects).filter(url=source)

    def query(self, constraint, sortby=None, typenames=None, maxrecords=10, startposition=0):
        ''' Query records from underlying repository '''

        # run the raw query and get total
        if 'where' in constraint:  # GetRecords with constraint
            query = self._get_repo_filter(Layer.objects).extra(where=[constraint['where']], params=constraint['values'])

        else:  # GetRecords sans constraint
            query = self._get_repo_filter(Layer.objects)

        total = query.count()

        # apply sorting, limit and offset
        if sortby is not None:
            if 'spatial' in sortby and sortby['spatial']:  # spatial sort
                desc = False
                if sortby['order'] == 'DESC':
                    desc = True
                query = query.all()
                return [str(total),
                        sorted(query,
                               key=lambda x: float(util.get_geometry_area(getattr(x, sortby['propertyname']))),
                               reverse=desc,
                               )[startposition:startposition+int(maxrecords)]]
            else:
                if sortby['order'] == 'DESC':
                    pname = '-%s' % sortby['propertyname']
                else:
                    pname = sortby['propertyname']
                return [str(total),
                        query.order_by(pname)[startposition:startposition+int(maxrecords)]]
        else:  # no sort
            return [str(total), query.all()[startposition:startposition+int(maxrecords)]]

    def insert(self, resourcetype, source, insert_date=None):
        ''' Insert a record into the repository '''

        caller = inspect.stack()[1][3]

        if caller == 'transaction':  # insert of Layer
            hhclass = 'Layer'
            source = resourcetype
            resourcetype = resourcetype.csw_schema
        else:  # insert of service
            hhclass = 'Service'
            if resourcetype not in HYPERMAP_SERVICE_TYPES.keys():
                raise RuntimeError('Unsupported Service Type')

        return self._insert_or_update(resourcetype, source, mode='insert', hhclass=hhclass)

    def _insert_or_update(self, resourcetype, source, mode='insert', hhclass='Service'):
        ''' Insert or update a record in the repository '''

        keywords = []

        if self.filter is not None:
            catalog = Catalog.objects.get(id=int(self.filter.split()[-1]))

        try:
            if hhclass == 'Layer':
                # TODO: better way of figuring out duplicates
                match = Layer.objects.filter(name=source.name,
                                             title=source.title,
                                             abstract=source.abstract,
                                             is_monitored=False)
                matches = match.all()
                if matches:
                    if mode == 'insert':
                        raise RuntimeError('HHypermap error: Layer %d \'%s\' already exists' % (
                                           matches[0].id, source.title))
                    elif mode == 'update':
                        match.update(
                            name=source.name,
                            title=source.title,
                            abstract=source.abstract,
                            is_monitored=False,
                            xml=source.xml,
                            wkt_geometry=source.wkt_geometry,
                            anytext=util.get_anytext([source.title, source.abstract, source.keywords_csv])
                        )

                service = get_service(source.xml)
                res, keywords = create_layer_from_metadata_xml(resourcetype, source.xml,
                                                               monitor=False, service=service,
                                                               catalog=catalog)

                res.save()

                LOGGER.debug('Indexing layer with id %s on search engine' % res.uuid)
                index_layer(res, use_cache=True)

            else:
                if resourcetype == 'http://www.opengis.net/cat/csw/2.0.2':
                    res = Endpoint(url=source, catalog=catalog)
                else:
                    res = Service(type=HYPERMAP_SERVICE_TYPES[resourcetype], url=source, catalog=catalog)

                res.save()

            if keywords:
                for kw in keywords:
                    res.keywords.add(kw)
        except Exception as err:
            raise RuntimeError('HHypermap error: %s' % err)

        # return a list of ids that were inserted or updated
        ids = []

        if hhclass == 'Layer':
            ids.append({'identifier': res.uuid, 'title': res.title})
        else:
            if resourcetype == 'http://www.opengis.net/cat/csw/2.0.2':
                for res in Endpoint.objects.filter(url=source).all():
                    ids.append({'identifier': res.uuid, 'title': res.url})
            else:
                for res in Service.objects.filter(url=source).all():
                    ids.append({'identifier': res.uuid, 'title': res.title})

        return ids

    def delete(self, constraint):
        ''' Delete a record from the repository '''

        results = self._get_repo_filter(Service.objects).extra(where=[constraint['where']],
                                                               params=constraint['values']).all()
        deleted = len(results)
        results.delete()
        return deleted

    def _get_repo_filter(self, query):
        ''' Apply repository wide side filter / mask query '''
        if self.filter is not None:
            return query.extra(where=[self.filter])
        return query
<|endoftext|>


<commit_before>from pathlib import Path

import psutil
from i3pystatus import Status

def get_cpu_temp_file(dir_path, filename, match):
    """Try to find which CPU file is the correct one"""
    dir = Path(dir_path)
    for p in dir.glob(filename):
        with open(p / "type") as f:
            if f.readline().rstrip() == match:
                return p / "temp"
    # failed to find a match, returns the first one
    return sorted(dir.glob(filename))[0] / "temp"

def get_mounted_block_devices(excludes=[]):
    """Find all mounted devices in the system"""
    result = []
    for disk in psutil.disk_partitions():
        if disk.mountpoint not in excludes:
            result.append(disk.mountpoint)
    return result

# create i3pystatus instance
status = Status()

# show clock
status.register(
    "clock",
    format=[" %H:%M:%S", " %a %d/%m"],
    on_leftclick="scroll_format",
)

# show current keyboard layout
status.register(
    "shell",
    command="setxkbmap -query | awk -F ': *' '/layout/ { print toupper($2) }'",
    interval=5,
    format=" {output}",
)

# show/change volume using PA
status.register(
    "pulseaudio",
    format=" {volume}%",
    format_muted=" Mute",
    on_leftclick="pavucontrol",
)

# show/control screen brightness
status.register(
    "backlight",
    format=" {percentage}%",
    format_no_backlight=" ",
)

# show network speed
status.register(
    "network",
    format_up="[ {essid} \[{quality}%\]]  {bytes_recv}K  {bytes_sent}K",
    format_down=" {interface}",
    next_if_down=True,
    on_upscroll=None,
    on_downscroll=None,
)

# show battery status
status.register(
    "battery",
    format="{status}{percentage:.0f}%[ {remaining}]",
    alert_percentage=10,
    not_present_text="",
    levels={
        25: " ",
        50: " ",
        75: " ",
        90: " "
    },
    status={
        "CHR": " ",
        "DPL": " ",
        "FULL": " ",
    },
)

# show disk available space
mounted_block_devices = get_mounted_block_devices(excludes=["/boot", "/nix/store"])
for block_device in mounted_block_devices[::-1]:
    status.register(
        "disk",
        format=" " + block_device + " {avail:.1f}G",
        path=block_device,
    )

# show available memory
status.register(
    "mem",
    format=" {avail_mem}G",
    warn_percentage=70,
    alert_percentage=90,
    divisor=1024**3,
)

# show cpu usage
status.register(
    "load",
    format=" {avg1} {avg5}",
)

# show CPU temperature
cpu_temp_file = get_cpu_temp_file("/sys/class/thermal",
                                  "thermal_zone*",
                                  "x86_pkg_temp")
status.register(
    "temp",
    file=cpu_temp_file,
    format=" {temp:.0f}°C",
)

# show CPU clock
status.register(
    "cpu_freq",
    format=" {avgg}",
)

# show current music info
status.register(
    "now_playing",
    format="{status} {title}",
    status={
        'play': '',
        'pause': '',
        'stop': '',
    }
)

status.register(
    "window_title",
    max_width=79
)

status.run()
<commit_msg>i3: Add dpms module in i3pystatus<commit_after>from pathlib import Path

import psutil
from i3pystatus import Status

def get_cpu_temp_file(dir_path, filename, match):
    """Try to find which CPU file is the correct one"""
    dir = Path(dir_path)
    for p in dir.glob(filename):
        with open(p / "type") as f:
            if f.readline().rstrip() == match:
                return p / "temp"
    # failed to find a match, returns the first one
    return sorted(dir.glob(filename))[0] / "temp"

def get_mounted_block_devices(excludes=[]):
    """Find all mounted devices in the system"""
    result = []
    for disk in psutil.disk_partitions():
        if disk.mountpoint not in excludes:
            result.append(disk.mountpoint)
    return result

# create i3pystatus instance
status = Status()

# show clock
status.register(
    "clock",
    format=[" %H:%M:%S", " %a %d/%m"],
    on_leftclick="scroll_format",
)

# show current keyboard layout
status.register(
    "shell",
    command="setxkbmap -query | awk -F ': *' '/layout/ { print toupper($2) }'",
    interval=5,
    format=" {output}",
)

# show/change volume using PA
status.register(
    "pulseaudio",
    format=" {volume}%",
    format_muted=" Mute",
    on_leftclick="pavucontrol",
)

status.register(
    "dpms",
    format=" ",
    format_disabled=" ",
)

# show/control screen brightness
status.register(
    "backlight",
    format=" {percentage}%",
    format_no_backlight=" ",
)

# show network speed
status.register(
    "network",
    format_up="[ {essid} \[{quality}%\]]  {bytes_recv}K  {bytes_sent}K",
    format_down=" {interface}",
    next_if_down=True,
    on_upscroll=None,
    on_downscroll=None,
)

# show battery status
status.register(
    "battery",
    format="{status}{percentage:.0f}%[ {remaining}]",
    alert_percentage=10,
    not_present_text="",
    levels={
        25: " ",
        50: " ",
        75: " ",
        90: " "
    },
    status={
        "CHR":  " ",
        "DPL":  " ",
        "FULL": " ",
    },
)

# show disk available space
mounted_block_devices = get_mounted_block_devices(excludes=["/boot", "/nix/store"])
for block_device in mounted_block_devices[::-1]:
    status.register(
        "disk",
        format=" " + block_device + " {avail:.1f}G",
        path=block_device,
    )

# show available memory
status.register(
    "mem",
    format=" {avail_mem}G",
    warn_percentage=70,
    alert_percentage=90,
    divisor=1024**3,
)

# show cpu usage
status.register(
    "load",
    format=" {avg1} {avg5}",
)

# show CPU temperature
cpu_temp_file = get_cpu_temp_file("/sys/class/thermal",
                                  "thermal_zone*",
                                  "x86_pkg_temp")
status.register(
    "temp",
    file=cpu_temp_file,
    format=" {temp:.0f}°C",
)

# show CPU clock
status.register(
    "cpu_freq",
    format=" {avgg}",
)

# show current music info
status.register(
    "now_playing",
    format="{status} {title}",
    status={
        "play":  "",
        "pause": "",
        "stop":  "",
    },
)

status.register(
    "window_title",
    max_width=79,
)

status.run()
<|endoftext|>