Facilitate cache directory override

Author: u8sand

maayanlab_bioinformatics/harmonization/ncbi_genes.py

@@ -3,7 +3,7 @@
 from maayanlab_bioinformatics.utils import fetch_save_read
 
 @lru_cache()
-def ncbi_genes_fetch(organism='Mammalia/Homo_sapiens', filters=lambda ncbi: ncbi['type_of_gene']=='protein-coding'):
+def ncbi_genes_fetch(organism='Mammalia/Homo_sapiens', filters=lambda ncbi: ncbi['type_of_gene']=='protein-coding', cache_dir=''):
   ''' Fetch the current NCBI Human Gene Info database.
   See ftp://ftp.ncbi.nih.gov/gene/DATA/GENE_INFO/ for the directory/file of the organism of interest.
   '''
@@ -28,6 +28,7 @@ def supplement_dbXref_prefix_omitted(ids):
     'ftp://ftp.ncbi.nih.gov/gene/DATA/GENE_INFO/{}.gene_info.gz'.format(organism),
     '{}.gene_info.tsv'.format(organism),
     sep='\t',
+    cache_dir=cache_dir,
   )
   if filters and callable(filters):
     ncbi = ncbi[filters(ncbi)]
@@ -47,15 +48,15 @@ def supplement_dbXref_prefix_omitted(ids):
   return ncbi
 
 @lru_cache()
-def ncbi_genes_lookup(organism='Mammalia/Homo_sapiens', filters=lambda ncbi: ncbi['type_of_gene']=='protein-coding'):
+def ncbi_genes_lookup(organism='Mammalia/Homo_sapiens', filters=lambda ncbi: ncbi['type_of_gene']=='protein-coding', cache_dir=''):
   ''' Return a lookup dictionary with synonyms as the keys, and official symbols as the values
   Usage:
   ```python
   ncbi_lookup = ncbi_genes_lookup('Mammalia/Homo_sapiens')
   print(ncbi_lookup('STAT3')) # any alias will get converted into the official symbol
   ```
   '''
-  ncbi_genes = ncbi_genes_fetch(organism=organism, filters=filters)
+  ncbi_genes = ncbi_genes_fetch(organism=organism, filters=filters, cache_dir=cache_dir)
   synonyms, symbols = zip(*{
     (synonym, gene_info['Symbol'])
     for _, gene_info in ncbi_genes.iterrows()

maayanlab_bioinformatics/utils/fetch_save_read.py

@@ -1,12 +1,17 @@
 import os
 import pandas as pd
 
-def fetch_save_read(url, file, reader=pd.read_csv, sep=',', **kwargs):
+def fetch_save_read(url, file, cache_dir='', reader=pd.read_csv, sep=',', **kwargs):
   ''' Download file from {url}, save it to {file}, and subsequently read it with {reader} using pandas options on {**kwargs}.
   '''
-  if not os.path.exists(file):
-    if os.path.dirname(file):
-      os.makedirs(os.path.dirname(file), exist_ok=True)
-    df = reader(url, sep=sep, index_col=None)
-    df.to_csv(file, sep=sep, index=False)
-  return pd.read_csv(file, sep=sep, **kwargs)
+  if cache_dir or cache_dir == '':
+    path = file if cache_dir == '' else cache_dir + os.path.sep + file
+    if not os.path.exists(path):
+      if os.path.dirname(path):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+      df = reader(url, sep=sep, index_col=None)
+      df.to_csv(path, sep=sep, index=False)
+    return reader(path, sep=sep, **kwargs)
+  else:
+    return reader(url, sep=sep, **kwargs)
+