<commit_before>#!/usr/bin/python
import argparse
import sys
import os
import os.path
import io
import time
import traceback
import ctypes
import subprocess
from subprocess import Popen, PIPE
import shutil
from optparse import OptionParser
import requests
from requests_toolbelt import MultipartEncoder
import json
import magic
import bz2
import tarfile
import zipfile
import glob
import re

try:
    from cStringIO import StringIO
except:
    from StringIO import StringIO


from biokbase.workspace.client import Workspace



BUF_SIZE = 8*1024 # default HTTP LIB client buffer_size

# Base class for Transform service
class TransformBase:

    def __init__(self, args):
        self.shock_url = args.shock_url
        self.inobj_id = args.inobj_id
        self.sdir = args.sdir
        self.itmp = args.itmp
        if(hasattr(args, 'otmp')):
          self.otmp = args.otmp
        else:
          self.otmp = "output"
        self.token = os.environ.get('KB_AUTH_TOKEN')
        self.ssl_verify = True


    def upload_to_shock(self):
        if self.token is None:
            raise Exception("Unable to find token!")
        
        filePath = "{}/{}".format(self.sdir,self.otmp)
    
        #build the header
        header = dict()
        header["Authorization"] = "Oauth %s" % self.token

        dataFile = open(os.path.abspath(filePath))
        m = MultipartEncoder(fields={'upload': (os.path.split(filePath)[-1], dataFile)})
        header['Content-Type'] = m.content_type

        try:
            response = requests.post(self.shock_url + "/node", headers=header, data=m, allow_redirects=True, verify=self.ssl_verify)
            dataFile.close()
    
            if not response.ok:
                response.raise_for_status()

            result = response.json()

            if result['error']:
                raise Exception(result['error'][0])
            else:
                return result["data"]    
        except:
            dataFile.close()
            raise
    
    
    def download_shock_data(self) :
        if self.token is None:
            raise Exception("Unable to find token!")
        
        # TODO: Improve folder checking
        if not os.path.isdir(self.sdir):
            try:
                os.mkdir(self.sdir)
            except:
                raise
        
        src_list = self.inobj_id.split(',')
        
        header = dict()
        header["Authorization"] = "OAuth {0}".format(self.token)
        
        # set chunk size to 10MB
        chunkSize = 10 * 2**20
    
        
        for sid in src_list:
            surl = "";
            fileName = ""
            fileSize = 0
            # TODO: let's improve shock node detection using actually trying it
            if (sid.startswith('http') or sid.startswith('ftp')) and not (re.search(r'^http[s]?.*/node/[a-fA-F0-9\-]+\?.*', sid)):
              surl = sid
              fileName = sid.split('/')[-1].split('#')[0].split('?')[0]
              #TODO: add file size estimation code here
            else:
              m = re.search(r'^(http[s]?.*/node/[a-fA-F0-9\-]+)\?.*', sid)
              if m is not None:
                metadata = requests.get(m.group(1), headers=header, stream=True, verify=self.ssl_verify)
                md = metadata.json()
                fileName = md['data']['file']['name']
                fileSize = md['data']['file']['size']
                metadata.close()
                surl = sid;
              else:
                metadata = requests.get("{0}/node/{1}?verbosity=metadata".format(self.shock_url, sid), headers=header, stream=True, verify=self.ssl_verify)
                md = metadata.json()
                fileName = md['data']['file']['name']
                fileSize = md['data']['file']['size']
                metadata.close()
                surl = "{0}/node/{1}?download_raw".format(self.shock_url, sid)

            data = requests.get(surl, headers=header, stream=True, verify=self.ssl_verify)

            size = int(data.headers['content-length'])

            if(size > 0 and fileSize == 0):
              fileSize = size
            
            filePath = os.path.join(self.sdir, fileName)
            <commit_msg>removed unuseful header for non-shock node<commit_after>#!/usr/bin/python
import argparse
import sys
import os
import os.path
import io
import time
import traceback
import ctypes
import subprocess
from subprocess import Popen, PIPE
import shutil
from optparse import OptionParser
import requests
from requests_toolbelt import MultipartEncoder
import json
import magic
import bz2
import tarfile
import zipfile
import glob
import re

try:
    from cStringIO import StringIO
except:
    from StringIO import StringIO


from biokbase.workspace.client import Workspace



BUF_SIZE = 8*1024 # default HTTP LIB client buffer_size

# Base class for Transform service
class TransformBase:

    def __init__(self, args):
        self.shock_url = args.shock_url
        self.inobj_id = args.inobj_id
        self.sdir = args.sdir
        self.itmp = args.itmp
        if(hasattr(args, 'otmp')):
          self.otmp = args.otmp
        else:
          self.otmp = "output"
        self.token = os.environ.get('KB_AUTH_TOKEN')
        self.ssl_verify = True


    def upload_to_shock(self):
        if self.token is None:
            raise Exception("Unable to find token!")
        
        filePath = "{}/{}".format(self.sdir,self.otmp)
    
        #build the header
        header = dict()
        header["Authorization"] = "Oauth %s" % self.token

        dataFile = open(os.path.abspath(filePath))
        m = MultipartEncoder(fields={'upload': (os.path.split(filePath)[-1], dataFile)})
        header['Content-Type'] = m.content_type

        try:
            response = requests.post(self.shock_url + "/node", headers=header, data=m, allow_redirects=True, verify=self.ssl_verify)
            dataFile.close()
    
            if not response.ok:
                response.raise_for_status()

            result = response.json()

            if result['error']:
                raise Exception(result['error'][0])
            else:
                return result["data"]    
        except:
            dataFile.close()
            raise
    
    
    def download_shock_data(self) :
        if self.token is None:
            raise Exception("Unable to find token!")
        
        # TODO: Improve folder checking
        if not os.path.isdir(self.sdir):
            try:
                os.mkdir(self.sdir)
            except:
                raise
        
        src_list = self.inobj_id.split(',')
        
        header = dict()
        header["Authorization"] = "OAuth {0}".format(self.token)
        
        # set chunk size to 10MB
        chunkSize = 10 * 2**20
    
        
        for sid in src_list:
            surl = "";
            fileName = ""
            fileSize = 0
            # TODO: let's improve shock node detection using actually trying it
            if (sid.startswith('http') or sid.startswith('ftp')) and not (re.search(r'^http[s]?.*/node/[a-fA-F0-9\-]+\?.*', sid)):
              surl = sid
              fileName = sid.split('/')[-1].split('#')[0].split('?')[0]
              #TODO: add file size estimation code here
              # reset header here with user id and password
              header = dict()
            else:
              m = re.search(r'^(http[s]?.*/node/[a-fA-F0-9\-]+)\?.*', sid)
              if m is not None:
                metadata = requests.get(m.group(1), headers=header, stream=True, verify=self.ssl_verify)
                md = metadata.json()
                fileName = md['data']['file']['name']
                fileSize = md['data']['file']['size']
                metadata.close()
                surl = sid;
              else:
                metadata = requests.get("{0}/node/{1}?verbosity=metadata".format(self.shock_url, sid), headers=header, stream=True, verify=self.ssl_verify)
                md = metadata.json()
                fileName = md['data']['file']['name']
                fileSize = md['data']['file']['size']
                metadata.close()
                surl = "{0}/node/{1}?download_raw".format(self.shock_url, sid)

            data = requests.get(surl, headers=header, stream=True, verify=self.ssl_verify)

            size = int(data.headers['content-length'])

            if(size > 0 and fileSize == 0):
              fileSize = size
            
            filePath = os.path.join(self.sdir, fileName)
            
{'fraction_non_alphanumeric': 0.12019739161085653, 'fraction_numerical': 0.007284690400657972, 'mean_word_length': 1.866958571909734, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4798661', 'n_tokens_mistral': 2534, 'n_tokens_neox': 2424, 'n_words': 587}