start script options for migration of Compound

(start working on #18)

ippisite/ippidb/management/commands/import_v1_data.py

@@ -2,7 +2,8 @@ from django.core.management import BaseCommand, CommandError
 import mysql.connector
 
 from ippidb.models import Bibliography, Protein, Taxonomy, MolecularFunction, \
-    Domain, ProteinDomainBoundComplex, ProteinDomainPartnerComplex, Symmetry, Ppi, PpiComplex, Disease
+    Domain, ProteinDomainBoundComplex, ProteinDomainPartnerComplex, Symmetry, Ppi, PpiComplex, Disease, \
+    Compound
 
 class Command(BaseCommand):
 
@@ -44,6 +45,13 @@ class Command(BaseCommand):
             default=False,
             help='Flush and migrate ppis and complexes',
         )
+        parser.add_argument(
+            '--compound',
+            action='store_true',
+            dest='compound',
+            default=False,
+            help='Flush and migrate compounds',
+        )
         parser.add_argument(
             '--stoponfail',
             action='store_true',
@@ -203,4 +211,26 @@ select distinct protein.NumUniprot, domain.PfamNumAccession  , complexe.NbCopy,
                         self.stdout.write(self.style.ERROR('Failed inserting {} {}'.format(row[0], row[1])))
                 else:
                     self.stdout.write(self.style.SUCCESS('Successfully inserted {} {}'.format(row[0], row[1])))
-            
+        if options['compound']:
+            cursor.execute("""SELECT * FROM compound""")
+            rows = cursor.fetchall()
+            Compound.objects.all().delete()
+            self.stdout.write(self.style.SUCCESS('Successfully flushed compound table'))
+            for row in rows:
+                try:
+                    compound = Compound()
+                    compound.canonical_smile = row[1]
+                    compound.is_macrocycle = (row[4]=='Y')
+                    compound.aromatic_ratio = row[5]
+                    compound.balaban_index = row[6]
+                    compound.balaban_index = row[6]
+                    compound.save()
+                except Exception as e:
+                    if options['stoponfail']:
+                        import traceback
+                        self.stderr.write(traceback.format_exc())
+                        raise CommandError('Failed inserting {} {}'.format(row[1], row[2]))
+                    else:
+                        self.stdout.write(self.style.ERROR('Failed inserting {} {}'.format(row[1], row[2])))
+                else:
+                    self.stdout.write(self.style.SUCCESS('Successfully inserted {} {}'.format(row[1], row[2])))