projet_court/src/pdb.py

import sys
#import collections


# custom imports
from atom import *

class PDBFile:

    def getContent(self, filename):
        with open(filename) as f:
            return(f.readlines())

    def getHeader(self):
        #Metadata = collections.namedtuple("Metadata", ["header", "compound", "source", "author"])
        Metadata = {}
        for line in self.rawLines:
            # no need to continue if meta are complete
            if(len(Metadata) <4):
                if(line[0:10] == "HEADER    "):
                    Metadata['header']=line
                elif(line[0:10] == "COMPND   2"):
                    Metadata['compound']=line
                elif(line[0:10] == "SOURCE   2"):
                    Metadata['source']=line
                elif(line[0:10] == "AUTHOR    "):
                    Metadata['author']=line
            else:
                # if meta are complete, stop parsing
                break     
        return(Metadata)

    def getAtoms(self):
        self.ATOMS = []
        self.RESIDUES = []
        tempAtoms = []
        for line in self.rawLines:
            if line.startswith("ATOM" or "HETATM"):
                atom = Atom(ATOM_ID = int(line[6:11].strip()),
                            ATOM_NAME = line[12:16].strip(),
                            ALT_LOCAT = line[16:17].strip(),
                            RES_NAME = line[17:20].strip(),
                            CHAIN_ID = line[21:22].strip(),
                            RES_SEQ_NB = int(line[22:26].strip()),
                            RES_INSER_CODE = line[26:27].strip(),
                            COORD_X = float(line[30:38].strip()),
                            COORD_Y = float(line[38:46].strip()),
                            COORD_Z = float(line[46:54].strip()),
                            OCCUPANCY = float(line[54:60].strip()),
                            TEMP_FACT = float(line[60:66].strip()),
                            ELEM_SYMBOL = line[76:78].strip(),
                            ATOM_CHARGE = line[78:80].strip())
                self.ATOMS.append(atom)
                # get the current indice of atom
                i = self.ATOMS.index(atom)
                # if this is a brand new residue
                if(len(self.ATOMS)>1
                   and atom.RES_SEQ_NB != self.ATOMS[i-1].RES_SEQ_NB):
                    self.RESIDUES.append(Residue(tempAtoms))
                    tempAtoms=[]
                tempAtoms.append(atom)
        # last residue
        self.RESIDUES.append(Residue(tempAtoms))
       
    def __init__(self, filename):
        self.rawLines = self.getContent(filename)
        self.Metadata = self.getHeader()
        self.getAtoms()
        for elem in self.Metadata :    
            print(self.Metadata[elem], end="")


if __name__ == "__main__":
    if(len(sys.argv)<2):
        print("Not enough arguments! Run with --help to learn more about proper"
              "call structure and parameters.")
    else:
        pdbFile = PDBFile(sys.argv[1])
        print(pdbFile.RESIDUES[0].ATOMS["N"].COORD_X)