Daryl Fung
fix content type error
36b9958
raw
history blame contribute delete
No virus
3.22 kB
"""This module proivdes wrapper functionality for the imgur API"""
import dotenv, os
from PIL import Image
import asyncio
from aiohttp import ClientSession
from typing import Optional, Union, Tuple
from .utils import image_to_b64_string, bytes_to_image
dotenv.load_dotenv()
API_ENDPOINTS = {
'upload': 'https://api.imgur.com/3/upload/',
'download': 'http://i.imgur.com/',
'info': 'https://api.imgur.com/3/image/',
'delete': 'https://api.imgur.com/3/image/',
'auth': f'https://api.imgur.com/oauth2/token'
}
# get access and refresh token
async def get_tokens():
session = ClientSession()
r = await session.request(
method='post',
url=API_ENDPOINTS['auth'],
data={
'refresh_token': os.getenv("IS3_REFRESH_TOKEN"),
'client_id': os.getenv("IS3_CLIENT_ID"),
'client_secret': os.getenv("IS3_CLIENT_SECRET"),
'grant_type': 'refresh_token',
}
)
r = await r.json(content_type='application/json')
return r['access_token'], r['refresh_token']
ACCESS_TOKEN, REFRESH_TOKEN = asyncio.run(get_tokens())
# AUTH_HEADER = {'Authorization': f"Client-ID {os.getenv('IS3_CLIENT_ID')}"}
AUTH_HEADER = {'Authorization': f"Bearer {ACCESS_TOKEN}"}
class ImgurClient:
"""Class to interact with various API endpoints"""
def __init__(self, session: Optional[ClientSession] = None) -> None:
self._session = session or ClientSession()
async def __aenter__(self):
return self
async def __aexit__(self, *err):
await self._session.close()
async def _request(self, method: str, url: str, *args, **kwargs) -> Union[dict, bytes]:
"""Make a request with the specified method to the endpoint. All requests
should either return raw image data as bytes or other data as JSON"""
async with self._session.request(method, url, *args, **kwargs) as resp:
content_type = resp.content_type
if content_type == 'image/png':
return await resp.read()
elif content_type == 'application/json':
return (await resp.json())['data']
else:
raise RuntimeError(f'Unexpected response content-type "{content_type}"')
async def upload_image(self, img: Image.Image) -> Tuple[str, str]:
"""Upload an image and return img id and deletehash"""
data = image_to_b64_string(img)
r = await self._request(
method='post',
url=API_ENDPOINTS['upload'],
headers=AUTH_HEADER,
data={'image': data, 'type': 'base64'}
)
return r['id'], r['deletehash']
async def download_image(self, image_id: str) -> Image.Image:
"""Download the image and return the data as bytes."""
url = API_ENDPOINTS['download'] + image_id + '.png'
data = await self._request('get', url)
return bytes_to_image(data)
async def delete_image(self, deletehash: str) -> None:
"""Delete an image using a deletehash string"""
url = API_ENDPOINTS['delete'] + deletehash
await self._request('delete', url, headers=AUTH_HEADER)