Skip to content

Commit

Permalink
build(deps): update dependencies (#124)
Browse files Browse the repository at this point in the history 
* build(deps): update dependencies

* fix(deps): use range dependency for urllib3 - not compatible yet

* fix: add vars to test.yml

* fix: improve and fix types
  • Loading branch information
@ailinvenerus
ailinvenerus authored Nov 21, 2023
1 parent 7537c67 commit d6f2509
Show file tree
Hide file tree
Showing 5 changed files with 156 additions and 94 deletions.
11 changes: 6 additions & 5 deletions .github/workflows/test.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@ name: Test

on:
push:
branches:
branches:
- master
pull_request:
branches:
branches:
- master
types: [opened, synchronize, reopened]

Expand All @@ -20,17 +20,18 @@ jobs:
python-version: '3.10'
cache: 'pip'
- name: Install deps
run: pip install -r requirements.txt
run: pip install -r requirements.txt
- name: Run tests
run: python -m coverage run -m pytest
env:
CLIENT_ID: ${{ secrets.CLIENT_ID }}
USERNAME: ${{ secrets.USERNAME }}
USERNAME2: ${{ secrets.USERNAME2 }}
PASSWORD: ${{ secrets.PASSWORD }}
OTPCODE: ${{ secrets.OTPCODE }}
USER_POOL_ID: ${{ secrets.USER_POOL_ID }}
FEDERATED_POOL_ID: ${{ secrets.FEDERATED_POOL_ID }}
IDENTITY_POOL_ID: ${{ secrets.IDENTITY_POOL_ID }}
API_URL: ${{ secrets.API_URL }}
REGION: ${{ secrets.REGION }}
- name: Publish code coverage
uses: paambaati/[email protected]
env:
Expand Down
9 changes: 6 additions & 3 deletions requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ boto3==1.29.3
botocore==1.32.3
certifi==2023.11.17
charset-normalizer==3.3.2
colorama==0.4.6
coverage==7.3.2
docutils==0.20.1
idna==3.4
Expand All @@ -18,20 +19,22 @@ nh3==0.2.14
packaging==23.2
pkginfo==1.9.6
pluggy==1.3.0
Pygments==2.17.0
Pygments==2.17.1
pytest==7.4.3
pytest-cov==4.1.0
python-dateutil==2.8.2
python-dotenv==1.0.0
pywin32-ctypes==0.2.2
readme-renderer==42.0
requests==2.31.0
requests-toolbelt==1.0.0
rfc3986==2.0.0
rich==13.7.0
s3transfer==0.7.0
scribeauth==1.0.8
scribeauth==1.1.1
six==1.16.0
twine==4.0.2
typing==3.7.4.3
typing_extensions==4.8.0
urllib3==2.0.7
urllib3>=2.0.7,<=2.1.0
zipp==3.17.0
189 changes: 116 additions & 73 deletions scribemi/ScribeMi.py
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
import requests
import json
from io import BytesIO
from typing import BinaryIO, Optional, TypedDict, Union
from typing import BinaryIO, Union
from scribeauth import ScribeAuth
from aws_requests_auth.aws_auth import AWSRequestsAuth
from datetime import datetime
from typing_extensions import TypedDict, Optional


class Env(TypedDict):
API_URL: str
Expand All @@ -13,46 +15,61 @@ class Env(TypedDict):
CLIENT_ID: str
REGION: str

class MITask(TypedDict):

class MITaskBase(TypedDict):
jobid: str
client: str
status: str
submitted: int


class MITask(MITaskBase, total=False):
companyName: Optional[str]
clientFilename: Optional[str]
originalFilename: Optional[str]
clientModelFilename: Optional[str]
status: str
submitted: int
modelUrl: Optional[str]

class SubmitTaskParams(TypedDict):

class SubmitTaskParamsBase(TypedDict):
filetype: str


class SubmitTaskParams(SubmitTaskParamsBase, total=False):
filename: Optional[str]
companyname: Optional[str]


class UnauthenticatedException(Exception):
"""
Exception raised when current token:
- Is wrong
- Has expired and needs to be updated
"""

pass


class TaskNotFoundException(Exception):
"""
Exception raised when trying to access an unexistent task.
"""

pass


class InvalidFiletypeException(Exception):
"""
Exception raised when trying to upload a file with a wrong filetype.
Accepted values:
'pdf', 'xlsx', 'xls', 'xlsm', 'doc', 'docx', 'ppt', 'pptx'
"""

pass


class MI:
def __init__(self, env):
"""
Expand All @@ -64,124 +81,150 @@ def __init__(self, env):
api_key -- The api key for the application.
"""
self.env = env
self.auth_client = ScribeAuth({
'client_id': env['CLIENT_ID'],
'user_pool_id': env['USER_POOL_ID'],
'identity_pool_id': env['IDENTITY_POOL_ID']
})
self.auth_client = ScribeAuth(
{
"client_id": env["CLIENT_ID"],
"user_pool_id": env["USER_POOL_ID"],
"identity_pool_id": env["IDENTITY_POOL_ID"],
}
)
self.tokens = None
self.user_id = None
self.request_auth = None

def authenticate(self, param):
self.tokens = self.auth_client.get_tokens(**param)
self.user_id = self.auth_client.get_federated_id(self.tokens['id_token'])
self.credentials = self.auth_client.get_federated_credentials(self.user_id, self.tokens['id_token'])
host = self.env['API_URL'].split('/')[0]
self.request_auth = AWSRequestsAuth(
aws_access_key=self.credentials['AccessKeyId'],
aws_secret_access_key=self.credentials['SecretKey'],
aws_token=self.credentials['SessionToken'],
aws_host=host,
aws_region=self.env['REGION'],
aws_service='execute-api'
)
id_token = self.tokens.get("id_token")
if id_token != None:
self.user_id = self.auth_client.get_federated_id(id_token)
self.credentials = self.auth_client.get_federated_credentials(
self.user_id, id_token
)
host = self.env["API_URL"].split("/")[0]
self.request_auth = AWSRequestsAuth(
aws_access_key=self.credentials["AccessKeyId"],
aws_secret_access_key=self.credentials["SecretKey"],
aws_token=self.credentials["SessionToken"],
aws_host=host,
aws_region=self.env["REGION"],
aws_service="execute-api",
)
else:
raise UnauthenticatedException("Authentication failed")

def reauthenticate(self):
if self.tokens == None or self.user_id == None:
raise UnauthenticatedException('Must authenticate before reauthenticating')
self.tokens = self.auth_client.get_tokens(refresh_token=self.tokens['refresh_token'])
self.credentials = self.auth_client.get_federated_credentials(self.user_id, self.tokens['id_token'])
host = self.env['API_URL'].split('/')[0]
self.request_auth = AWSRequestsAuth(
aws_access_key=self.credentials['AccessKeyId'],
aws_secret_access_key=self.credentials['SecretKey'],
aws_token=self.credentials['SessionToken'],
aws_host=host,
aws_region=self.env['REGION'],
aws_service='execute-api'
)
raise UnauthenticatedException("Must authenticate before reauthenticating")
refresh_token = self.tokens.get("refresh_token")
if refresh_token != None:
self.tokens = self.auth_client.get_tokens(refresh_token=refresh_token)
id_token = self.tokens.get("id_token")
if id_token != None:
self.credentials = self.auth_client.get_federated_credentials(
self.user_id, id_token
)
host = self.env["API_URL"].split("/")[0]
self.request_auth = AWSRequestsAuth(
aws_access_key=self.credentials["AccessKeyId"],
aws_secret_access_key=self.credentials["SecretKey"],
aws_token=self.credentials["SessionToken"],
aws_host=host,
aws_region=self.env["REGION"],
aws_service="execute-api",
)
else:
raise UnauthenticatedException("Authentication failed")
else:
raise UnauthenticatedException("Authentication failed")

def call_endpoint(self, method, path, data=None, params=None):
if self.request_auth == None:
raise UnauthenticatedException('Not authenticated')
if self.credentials['Expiration'] < datetime.now(self.credentials['Expiration'].tzinfo):
raise UnauthenticatedException("Not authenticated")
if self.credentials["Expiration"] < datetime.now(
self.credentials["Expiration"].tzinfo
):
self.reauthenticate()
res = requests.request(
method=method,
url='https://{host}{path}'.format(host=self.env['API_URL'], path=path),
url="https://{host}{path}".format(host=self.env["API_URL"], path=path),
params=params,
json=data,
auth=self.request_auth,
)
if res.status_code == 200:
return json.loads(res.text)
elif res.status_code == 401 or res.status_code == 403:
raise UnauthenticatedException('Authentication failed ({})'.format(res.status_code))
raise UnauthenticatedException(
"Authentication failed ({})".format(res.status_code)
)
elif res.status_code == 404:
raise TaskNotFoundException('Not found')
raise TaskNotFoundException("Not found")
else:
raise Exception('Unexpected error ({})'.format(res.status_code))
raise Exception("Unexpected error ({})".format(res.status_code))

def list_tasks(self, companyName=None) -> list[MITask]:
params = {
'includePresigned': True
}
params = {"includePresigned": True}
if companyName != None:
params['company'] = companyName
return self.call_endpoint('GET', '/tasks', params=params).get('tasks')
params["company"] = companyName
return self.call_endpoint("GET", "/tasks", params=params).get("tasks")

def get_task(self, jobid: str) -> MITask:
return self.call_endpoint('GET', '/tasks/{}'.format(jobid))
return self.call_endpoint("GET", "/tasks/{}".format(jobid))

def fetch_model(self, task: MITask):
if task.modelUrl == None:
raise Exception('Cannot load model for task {}: model is not ready for export'.format(task.jobid))
res = requests.get(task.modelUrl)
if res.status == 200:
modelUrl = task.get("modelUrl")
if modelUrl == None:
raise Exception(
"Cannot load model for task {}: model is not ready for export".format(
task.get("jobid")
)
)
res = requests.get(modelUrl)
if res.status_code == 200:
return json.loads(res.text)
elif res.status_code == 401 or res.status_code == 403:
raise UnauthenticatedException(
'{} Authentication failed (possibly due to timeout: try calling get_task immediately before fetch_model)'
.format(res.status_code)
"{} Authentication failed (possibly due to timeout: try calling get_task immediately before fetch_model)".format(
res.status_code
)
)
else:
raise Exception('Unexpected error ({})'.format(res.status_code))
raise Exception("Unexpected error ({})".format(res.status_code))

def consolidate_tasks(self, tasks: list[MITask]):
jobids = list(
map(
lambda task: task.jobid,
tasks
)
)
jobids_param = ';'.join(jobids)
res = self.call_endpoint('GET', '/fundportfolio?jobids={}'.format(jobids_param))
jobids = list(map(lambda task: task.get("jobid"), tasks))
jobids_param = ";".join(jobids)
res = self.call_endpoint("GET", "/fundportfolio?jobids={}".format(jobids_param))
return res.model

def submit_task(self, file_or_filename: Union[str, BytesIO, BinaryIO], params: SubmitTaskParams):
filetype_list = ['pdf', 'xlsx', 'xls', 'xlsm', 'doc', 'docx', 'ppt', 'pptx']
if params.get('filetype') not in filetype_list:
raise InvalidFiletypeException("Invalid filetype. Accepted values: 'pdf', 'xlsx', 'xls', 'xlsm', 'doc', 'docx', 'ppt', 'pptx'.")
def submit_task(
self, file_or_filename: Union[str, BytesIO, BinaryIO], params: SubmitTaskParams
):
filetype_list = ["pdf", "xlsx", "xls", "xlsm", "doc", "docx", "ppt", "pptx"]
if params.get("filetype") not in filetype_list:
raise InvalidFiletypeException(
"Invalid filetype. Accepted values: 'pdf', 'xlsx', 'xls', 'xlsm', 'doc', 'docx', 'ppt', 'pptx'."
)

if isinstance(file_or_filename, str) and params.get('filename') == None:
params['filename'] = file_or_filename
if isinstance(file_or_filename, str) and params.get("filename") == None:
params["filename"] = file_or_filename

post_res = self.call_endpoint('POST', '/tasks', params)
put_url = post_res['url']
post_res = self.call_endpoint("POST", "/tasks", params)
put_url = post_res["url"]

if isinstance(file_or_filename, str):
with open(file_or_filename, 'rb') as file:
with open(file_or_filename, "rb") as file:
upload_file(file, put_url)
else:
return upload_file(file_or_filename, put_url)

return post_res['jobid']
return post_res["jobid"]

def delete_task(self, task: MITask):
return self.call_endpoint('DELETE', '/tasks/{}'.format(task['jobid']))
return self.call_endpoint("DELETE", "/tasks/{}".format(task["jobid"]))


def upload_file(file, url):
res = requests.put(url, data=file)
if res.status_code != 200:
raise Exception('Error uploading file: {}'.format(res.status_code))
raise Exception("Error uploading file: {}".format(res.status_code))
2 changes: 1 addition & 1 deletion scribemi/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1 +1 @@
from .ScribeMi import MI
from .ScribeMi import MI
Loading

0 comments on commit d6f2509

Please sign in to comment.