Merge pull request #83 from sbassi/subjectsDB

Support for subjects in SQLiteDB

Author: mpacer

allofplos/corpus/corpus_class.py

@@ -30,7 +30,7 @@ def __repr__(self):
     def iter_file_doi(self):
         """Generator that returns filename, doi tuples for every file in the corpus.
 
-        Used to generate both DOI and file generators for the corpus. 
+        Used to generate both DOI and file generators for the corpus.
         """
         return ((file_, filename_to_doi(file_))
                 for file_ in os.listdir(self.directory)

allofplos/makedb.py

@@ -17,15 +17,11 @@
     DateTimeField, BooleanField, IntegerField, IntegrityError
 from playhouse.sqlite_ext import SqliteExtDatabase
 
+from .corpus import Corpus
 from .transformations import filename_to_doi, convert_country
+from . import starterdir
 from .article_class import Article
 
-
-# TODO: this may need to be updated to take into account the new get_corpus_dir() logic.
-# It is not clear, since this is a relative path from the package how this should work without
-# diving into the code more thoroughly 
-corpusdir = 'allofplos/allofplos_xml'
-
 journal_title_dict = {
     'PLOS ONE': 'PLOS ONE',
     'PLOS GENETICS': 'PLOS Genetics',
@@ -83,6 +79,9 @@ class Country(BaseModel):
 class Affiliations(BaseModel):
     affiliations = CharField(unique=True)
 
+class Subjects(BaseModel):
+    subjects = CharField(unique=True)
+
 class CorrespondingAuthor(BaseModel):
     corr_author_email = CharField(unique=True)
     tld = TextField(null=True)
@@ -105,26 +104,25 @@ class PLOSArticle(BaseModel):
     word_count = IntegerField()
     JATS_type = ForeignKeyField(JATSType, related_name='jats')
 
+class SubjectsPLOSArticle(BaseModel):
+    subject = ForeignKeyField(Subjects)
+    article = ForeignKeyField(PLOSArticle)
+
 class CoAuthorPLOSArticle(BaseModel):
     corr_author = ForeignKeyField(CorrespondingAuthor)
     article = ForeignKeyField(PLOSArticle)
 
 db.connect()
-db.create_tables([Journal, PLOSArticle, ArticleType,
-                      CoAuthorPLOSArticle, CorrespondingAuthor,
-                      JATSType, Affiliations, Country])
+db.create_tables([Journal, PLOSArticle, ArticleType, CoAuthorPLOSArticle,
+                  CorrespondingAuthor, JATSType, Affiliations, Country,
+                  SubjectsPLOSArticle, Subjects])
 
-if args.starter:
-    allfiles = os.listdir('starter_corpus')
-else:
-    allfiles = os.listdir(corpusdir)
-if args.random:
-    randomfiles = random.sample(allfiles, args.random)
-    max_value = len(randomfiles)
-else:
-    max_value = len(allfiles)
 
-for i, file_ in enumerate(tqdm(randomfiles if args.random else allfiles)):
+corpus_dir = starterdir if args.starter else None
+allfiles = Corpus(corpus_dir).files
+files = random.sample(allfiles, args.random) if args.random else allfiles
+
+for file_ in tqdm(files):
     doi = filename_to_doi(file_)
     article = Article(doi)
     journal_name = journal_title_dict[article.journal.upper()]
@@ -155,6 +153,24 @@ class CoAuthorPLOSArticle(BaseModel):
         created_date = article.pubdate,
         word_count=article.word_count,
         JATS_type = j_type)
+    # Get subject information
+    taxonomy_set = set()
+    taxonomy = article.taxonomy
+    for values in taxonomy.values():
+        for value in values:
+            for taxon in value:
+                taxonomy_set.add(taxon)
+    for taxon in taxonomy_set:
+        with db.atomic() as atomic:
+            try:
+                subject = Subjects.create(subjects = taxon)
+            except (sqlite3.IntegrityError, IntegrityError):
+                db.rollback()
+                subject = Subjects.get(Subjects.subjects == taxon)
+        SubjectsPLOSArticle.create(
+                subject = subject,
+                article = p_art
+            )
     if article.authors:
         iterable_authors = article.authors
     else: