<commit_before>import os

import requests
import globus_sdk
from tqdm import tqdm

# MDF Utils
from mdf_indexers.utils import auth
from mdf_indexers.utils.gmeta_utils import gmeta_pop
from mdf_indexers.utils.globus_utils import get_local_ep

HTTP_NUM_LIMIT = 10


def build_source_list(sources=[], match_all=False):
    join_term = "AND" if match_all else "OR"
    sources = ["mdf.source_name:"+s for s in sources]
    return "("+ (' ' + join_term + ' ').join(sources) + ")"

def get_content_block(res):
    #TODO Check if GlobusHTTP ressonse
    tf = [ r['content'][0] for r in res['gmeta']]
    return tf

def get_file_http(loc, save_files=True):
    response = requests.get(loc.get('from'), stream=True)
    if not response.ok: return False
    
    # Throw an error for bad status codes
    response.raise_for_status()

    with open(loc.get('to'), 'wb') as handle:
        for block in response.iter_content(1024):
            handle.write(block)
    #return response.ok         
    
def get_files(locs=[],by_http=True, by_globus=False, n_workers=1):
    tasks = res['gmeta']
    #pbar = tqdm.tqdm(total=len(tasks)/n_workers)

    if by_globus is True:
        #perform Globus trans
        pass
    elif by_http is True:
        if n_workers>1:
            mp = Pool(n_workers)
            mdf_data = mp.map(get_file_http, locs)
            mp.close()
            mp.join()
        else:
            for loc in locs:
                get_file_http(loc)
    else:
        pass
    

class Forge:
    index = "mdf"
    services = ["mdf", "transfer", "search_ingest"]
    app_name = "MDF Forge"

    def __init__(self, data={}):
        self.index = data.get('index', self.index)
        self.services = data.get('services', self.services)
        self.local_ep = data.get("local_ep", None)

        clients = auth.login(credentials={
                                "app_name": self.app_name,
                                "services": self.services,
                                "index": self.index})
        self.search_client = clients["search"]
        self.transfer_client = clients["transfer"]
        self.mdf_authorizer = clients["mdf"]
    

    def search(self, q, raw=False, advanced=False):
        if not advanced:
            res = self.search_client.search(q)
        else:
            query = {
                "q": q,
                "advanced": True,
                "limit": 9999
                }
            res = self.search_client.structured_search(query)
        return res if raw else gmeta_pop(res)


    def search_by_elements(self, *elements, sources=[], match_all=False, raw=False):
        elements = list(elements)
        q_sources = (build_source_list(sources) + " AND ") if sources else ""
        if match_all:
            q_elements = " AND ".join(["mdf.elements:"+elem for elem in elements])
        else:
            q_elements = "mdf.elements:" + ','.join(elements)
        q = {
            "q": (q_sources +
                  "mdf.resource_type:record AND " +
                  q_elements),
            "advanced": True,
            "limit":200
        }            
        
        res = self.search_client.structured_search(q)
        
        return res if raw else gmeta_pop(res)


    def get_http(self, results, dest=".", preserve_dir=False, verbose=True):
        if type(results) is globus_sdk.GlobusHTTPResponse:
            results = gmeta_pop(results)
        if len(results) > HTTP_NUM_LIMIT:
            return {
                "success": False,
                "message": "Too many results supplied. Use get_globus() for fetching more than " + str(HTTP_NUM_LIMIT) + " entries."
                }
        for res in tqdm(results, desc="Fetching files", disable= not verbose):
            for key in tqdm(res["mdf"]["links"].keys(), desc="Fetching files", disable=True):  # not verbose):<commit_msg>minor changes to search_by_elements. Now expects elements as a list and also allows a custom limit to be set with the limit arg<commit_after>import os

import requests
import globus_sdk
from tqdm import tqdm

# MDF Utils
from mdf_indexers.utils import auth
from mdf_indexers.utils.gmeta_utils import gmeta_pop
from mdf_indexers.utils.globus_utils import get_local_ep

HTTP_NUM_LIMIT = 10


def build_source_list(sources=[], match_all=False):
    join_term = "AND" if match_all else "OR"
    sources = ["mdf.source_name:"+s for s in sources]
    return "("+ (' ' + join_term + ' ').join(sources) + ")"

def get_content_block(res):
    #TODO Check if GlobusHTTP ressonse
    tf = [ r['content'][0] for r in res['gmeta']]
    return tf

def get_file_http(loc, save_files=True):
    response = requests.get(loc.get('from'), stream=True)
    if not response.ok: return False
    
    # Throw an error for bad status codes
    response.raise_for_status()

    with open(loc.get('to'), 'wb') as handle:
        for block in response.iter_content(1024):
            handle.write(block)
    #return response.ok         
    
def get_files(locs=[],by_http=True, by_globus=False, n_workers=1):
    tasks = res['gmeta']
    #pbar = tqdm.tqdm(total=len(tasks)/n_workers)

    if by_globus is True:
        #perform Globus trans
        pass
    elif by_http is True:
        if n_workers>1:
            mp = Pool(n_workers)
            mdf_data = mp.map(get_file_http, locs)
            mp.close()
            mp.join()
        else:
            for loc in locs:
                get_file_http(loc)
    else:
        pass
    

class Forge:
    index = "mdf"
    services = ["mdf", "transfer", "search_ingest"]
    app_name = "MDF Forge"

    def __init__(self, data={}):
        self.index = data.get('index', self.index)
        self.services = data.get('services', self.services)
        self.local_ep = data.get("local_ep", None)

        clients = auth.login(credentials={
                                "app_name": self.app_name,
                                "services": self.services,
                                "index": self.index})
        self.search_client = clients["search"]
        self.transfer_client = clients["transfer"]
        self.mdf_authorizer = clients["mdf"]
    

    def search(self, q, raw=False, advanced=False):
        if not advanced:
            res = self.search_client.search(q)
        else:
            query = {
                "q": q,
                "advanced": True,
                "limit": 9999
                }
            res = self.search_client.structured_search(query)
        return res if raw else gmeta_pop(res)


    def search_by_elements(self, elements=[], sources=[],limit=200, match_all=False, raw=False):
        q_sources = (build_source_list(sources) + " AND ") if sources else ""
        if match_all:
            q_elements = " AND ".join(["mdf.elements:"+elem for elem in elements])
        else:
            q_elements = "mdf.elements:" + ','.join(elements)
        q = {
            "q": (q_sources +
                  "mdf.resource_type:record AND " +
                  q_elements),
            "advanced": True,
            "limit":limit
        }            
        
        res = self.search_client.structured_search(q)
        
        return res if raw else gmeta_pop(res)


    def get_http(self, results, dest=".", preserve_dir=False, verbose=True):
        if type(results) is globus_sdk.GlobusHTTPResponse:
            results = gmeta_pop(results)
        if len(results) > HTTP_NUM_LIMIT:
            return {
                "success": False,
                "message": "Too many results supplied. Use get_globus() for fetching more than " + str(HTTP_NUM_LIMIT) + " entries."
                }
        for res in tqdm(results, desc="Fetching files", disable= not verbose):
            for key in tqdm(res["mdf"]["links"].keys(), desc="Fetching files", disable=True):  # not verbose):
{'fraction_non_alphanumeric': 0.11383985674085444, 'fraction_numerical': 0.004093118444614991, 'mean_word_length': 2.118867171918628, 'pattern_counts': {'":': 22, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6548933', 'n_tokens_mistral': 2451, 'n_tokens_neox': 2387, 'n_words': 585}