# $Id$
#
# Copyright (C) 2007-2010 Greg Landrum
# All Rights Reserved
#
from rdkit import Chem

class PropertyMol(Chem.Mol):
  """ allows rdkit molecules to be pickled with their properties saved.

   >>> import cPickle
   >>> m = Chem.MolFromMolFile('test_data/benzene.mol')
   >>> m.GetProp('_Name')
   'benzene.mol'

   by default pickling removes properties:
   >>> m2 = cPickle.loads(cPickle.dumps(m))
   >>> m2.HasProp('_Name')
   0

   Property mols solve this:
   >>> pm = PropertyMol(m)
   >>> pm.GetProp('_Name')
   'benzene.mol'
   >>> pm.SetProp('MyProp','foo')
   >>> pm.HasProp('MyProp')
   1

   >>> pm2 = cPickle.loads(cPickle.dumps(pm))
   >>> Chem.MolToSmiles(pm2)
   'c1ccccc1'
   >>> pm2.GetProp('_Name')
   'benzene.mol'
   >>> pm2.HasProp('MyProp')
   1
   >>> pm2.GetProp('MyProp')
   'foo'
   >>> pm2.HasProp('MissingProp')
   0

   Property mols are a bit more permissive about the types
   of property values:
   >>> pm.SetProp('IntVal',1)

   That wouldn't work with a standard mol:
   >>> m.SetProp('IntVal',1)
   Traceback (most recent call last):
     ...
   ArgumentError: Python argument types in
       Mol.SetProp(Mol, str, int)
   did not match C++ signature:
     ...

   but the Property mols still convert all values to strings before storing:
   >>> pm.GetProp('IntVal')
   '1'
   
   This is a test for sf.net issue 2880943: make sure properties end up in SD files:
   >>> import tempfile,os
   >>> fn = tempfile.mktemp('.sdf')
   >>> w = Chem.SDWriter(fn)
   >>> w.write(pm)
   >>> w=None
   >>> txt = file(fn,'r').read()
   >>> '<IntVal>' in txt
   True
   >>> try:
   ...   os.unlink(fn)
   ... except:
   ...   pass

   The next level of that bug: does writing a *depickled* propertymol
   to an SD file include properties:
   >>> fn = tempfile.mktemp('.sdf')
   >>> w = Chem.SDWriter(fn)
   >>> pm = cPickle.loads(cPickle.dumps(pm))
   >>> w.write(pm)
   >>> w=None
   >>> txt = file(fn,'r').read()
   >>> '<IntVal>' in txt
   True
   >>> try:
   ...   os.unlink(fn)
   ... except:
   ...   pass


   
  """
  __getstate_manages_dict__=True
  def __init__(self,mol):
    if not isinstance(mol,Chem.Mol): return
    Chem.Mol.__init__(self,mol)
    for pn in mol.GetPropNames(includePrivate=True):
      self.SetProp(pn,mol.GetProp(pn))
  def SetProp(self,nm,val):
    Chem.Mol.SetProp(self,nm,str(val))
  def __getstate__(self):
    pDict={}
    for pn in self.GetPropNames(includePrivate=True):
      pDict[pn] = self.GetProp(pn)
    return {'pkl':self.ToBinary(),
            'propD':pDict}
  def __setstate__(self,stateD):
    Chem.Mol.__init__(self,stateD['pkl'])
    for prop,val in stateD['propD'].iteritems():
      self.SetProp(prop,val)

    
#------------------------------------
#
#  doctest boilerplate
#
def _test():
  import doctest,sys
  return doctest.testmod(sys.modules["__main__"],optionflags=doctest.ELLIPSIS)

if __name__ == '__main__':
  import sys
  failed,tried = _test()
  sys.exit(failed)