<commit_before>from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .arrow_uploader import ArrowUploader

import logging, pyarrow as pa, requests, sys
from functools import lru_cache
from weakref import WeakKeyDictionary

logger = logging.getLogger('ArrowFileUploader')


# WrappedTable -> {'file_id': str, 'output': dict}
DF_TO_FILE_ID_CACHE : WeakKeyDictionary = WeakKeyDictionary()
"""
NOTE: Will switch to pa.Table -> ... when RAPIDS upgrades from pyarrow, 
     which adds weakref support
"""

class ArrowFileUploader():
    """
        Implement file API with focus on Arrow support

        Memoization in this class is based on reference equality, while plotter is based on hash.
        That means the plotter resolves different-identity value matches, so by the time ArrowFileUploader compares,
        identities are unified for faster reference-based checks.

        Example: Upload files with per-session memoization
            uploader : ArrowUploader
            arr : pa.Table
            afu = ArrowFileUploader(uploader)

            file1_id = afu.create_and_post_file(arr)[0]
            file2_id = afu.create_and_post_file(arr)[0]

            assert file1_id == file2_id # memoizes by default (memory-safe: weak refs)

        Example: Explicitly create a file and upload data for it
            uploader : ArrowUploader
            arr : pa.Table
            afu = ArrowFileUploader(uploader)

            file1_id = afu.create_file()
            afu.post_arrow(arr, file_id)

            file2_id = afu.create_file()
            afu.post_arrow(arr, file_id)

            assert file1_id != file2_id

    """

    uploader: 'ArrowUploader'

    def __init__(self, uploader: 'ArrowUploader'):
        self.uploader = uploader

    ###

    def create_file(self, file_opts: dict = {}) -> str:
        """
            Creates File and returns file_id str.
            
            Defauls:
              - file_type: 'arrow'

            See File REST API for file_opts

        """

        tok = self.uploader.token

        json_extended = {
            'file_type': 'arrow',
            'agent_name': 'pygraphistry',
            'agent_version': sys.modules['graphistry'].__version__,
            **file_opts
        }

        res = requests.post(
            self.uploader.server_base_path + '/api/v2/files/',
            verify=self.uploader.certificate_validation,
            headers={'Authorization': f'Bearer {tok}'},
            json=json_extended)

        try:            
            out = res.json()
            logger.debug('Server create file response: %s', out)
            if (not 'success' in out) or (not out['success']):
                raise Exception(out)
        except Exception as e:
            logger.error('Failed creating file: %s', res.text, exc_info=True)
            raise e
        
        self.dataset_id = out['data']['file_id']

        return out

    def post_arrow(self, arr: pa.Table, file_id: str, url_opts: str = 'erase=true') -> dict:
        """
            Upload new data to existing file id

            Default url_opts='erase=true' throws exceptions on parse errors and deletes upload.

            See File REST API for url_opts (file upload)
        """

        buf = self.uploader.arrow_to_buffer(arr)

        tok = self.uploader.token
        base_path = self.uploader.server_base_path

        url = f'{base_path}/api/v2/upload/files/{file_id}'
        if len(url_opts) > 0:
            url = f'{url}?{url_opts}'

        out = requests.post(
            url,
            verify=self.uploader.certificate_validation,
            headers={'Authorization': f'Bearer {tok}'},
            data=buf).json()
        
        if not out['success']:
            raise Exception(out)
            <commit_msg>fix(file upload): handle return type<commit_after>from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from .arrow_uploader import ArrowUploader

import logging, pyarrow as pa, requests, sys
from functools import lru_cache
from weakref import WeakKeyDictionary

logger = logging.getLogger('ArrowFileUploader')


# WrappedTable -> {'file_id': str, 'output': dict}
DF_TO_FILE_ID_CACHE : WeakKeyDictionary = WeakKeyDictionary()
"""
NOTE: Will switch to pa.Table -> ... when RAPIDS upgrades from pyarrow, 
     which adds weakref support
"""

class ArrowFileUploader():
    """
        Implement file API with focus on Arrow support

        Memoization in this class is based on reference equality, while plotter is based on hash.
        That means the plotter resolves different-identity value matches, so by the time ArrowFileUploader compares,
        identities are unified for faster reference-based checks.

        Example: Upload files with per-session memoization
            uploader : ArrowUploader
            arr : pa.Table
            afu = ArrowFileUploader(uploader)

            file1_id = afu.create_and_post_file(arr)[0]
            file2_id = afu.create_and_post_file(arr)[0]

            assert file1_id == file2_id # memoizes by default (memory-safe: weak refs)

        Example: Explicitly create a file and upload data for it
            uploader : ArrowUploader
            arr : pa.Table
            afu = ArrowFileUploader(uploader)

            file1_id = afu.create_file()
            afu.post_arrow(arr, file_id)

            file2_id = afu.create_file()
            afu.post_arrow(arr, file_id)

            assert file1_id != file2_id

    """

    uploader: 'ArrowUploader'

    def __init__(self, uploader: 'ArrowUploader'):
        self.uploader = uploader

    ###

    def create_file(self, file_opts: dict = {}) -> str:
        """
            Creates File and returns file_id str.
            
            Defauls:
              - file_type: 'arrow'

            See File REST API for file_opts

        """

        tok = self.uploader.token

        json_extended = {
            'file_type': 'arrow',
            'agent_name': 'pygraphistry',
            'agent_version': sys.modules['graphistry'].__version__,
            **file_opts
        }

        res = requests.post(
            self.uploader.server_base_path + '/api/v2/files/',
            verify=self.uploader.certificate_validation,
            headers={'Authorization': f'Bearer {tok}'},
            json=json_extended)

        try:            
            out = res.json()
            if res.status_code != requests.codes.ok:
                res.raise_for_status()
            if not out['is_valid']:
                if out['is_uploaded']:
                    raise RuntimeError("Uploaded file but cannot parse, see errors", out['errors'])
                else:
                    raise RuntimeError("Erased file upon failure, see errors", out['errors'])
            logger.debug('Server create file response: %s', out)
        except Exception as e:
            logger.error('Failed creating file: %s', res.text, exc_info=True)
            raise e
        
        return out['file_id']

    def post_arrow(self, arr: pa.Table, file_id: str, url_opts: str = 'erase=true') -> dict:
        """
            Upload new data to existing file id

            Default url_opts='erase=true' throws exceptions on parse errors and deletes upload.

            See File REST API for url_opts (file upload)
        """

        buf = self.uploader.arrow_to_buffer(arr)

        tok = self.uploader.token
        base_path = self.uploader.server_base_path

        url = f'{base_path}/api/v2/upload/files/{file_id}'
        if len(url_opts) > 0:
            url = f'{url}?{url_opts}'

        out = requests.post(
            url,
            verify=self.uploader.certificate_validation,
            headers={'Authorization': f'Bearer {tok}'},
            data=buf).json()
        
        if not out['success']:
            raise Exception(out)
            
{'fraction_non_alphanumeric': 0.09282162437842663, 'fraction_numerical': 0.0033150580135152363, 'mean_word_length': 2.0568978955572876, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, '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': '4645685', 'n_tokens_mistral': 2264, 'n_tokens_neox': 2176, 'n_words': 600}