#
# 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 os
   >>> from rdkit.six.moves import cPickle
   >>> from rdkit import RDConfig
   >>> m = Chem.MolFromMolFile(os.path.join(RDConfig.RDCodeDir, 'Chem', '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

   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 = open(fn,'r').read()
   >>> '<IntVal>' in txt
   True
   >>> try:
   ...   os.unlink(fn)
   ... except Exception:
   ...   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 = open(fn,'r').read()
   >>> '<IntVal>' in txt
   True
   >>> try:
   ...   os.unlink(fn)
   ... except Exception:
   ...   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'].items():
      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)
