Wheelie/generator/main.py
2021-02-24 14:58:26 +01:00

154 lines
5.4 KiB
Python

from argparse import ArgumentParser
from io import BytesIO
from json import dumps, loads
from ntpath import basename
from os import environ, listdir
from os.path import isfile, join
import botocore.exceptions as botoexeptions
from boto3 import client
from jinja2 import Environment, FileSystemLoader
if not environ.get("CI"):
from dotenv import load_dotenv
load_dotenv(dotenv_path="../.dev/.env")
mime = {"json": "application/json",
"whl": "application/x-wheel+zip", # Not a "real" type
"html": "text/html"
}
class Bucket():
def __init__(self) -> None:
self.client = client(
"s3",
region_name=environ.get("S3_REGION"),
aws_access_key_id=environ.get("S3_ACCESS_KEY"),
aws_secret_access_key=environ.get("S3_SECRET_KEY"),
endpoint_url=environ.get("S3_IS_MINIO", None)
)
self.bucket_name = environ.get("S3_BUCKET")
self._list()
def _list(self) -> list:
response = self.client.list_buckets()
buckets = []
for bucket in response["Buckets"]:
buckets.append(bucket)
return buckets
def download(self, bucket, filepath) -> bytes:
_file = BytesIO()
self.client.download_fileobj(bucket, filepath, _file)
_file.seek(0)
return _file.read()
def upload(self, bucket, filepath, file_):
m_type = mime[filepath.split(".")[-1]]
self.client.upload_fileobj(file_, bucket, filepath, ExtraArgs={
"ContentType": m_type, "ACL": "public-read", "CacheControl": "no-cache"})
class Index(Bucket):
def __init__(self, distribution: str, release: str, folder: str) -> None:
super().__init__()
self.distro = distribution
self.rel = release
self.path = folder
self.get_manifest()
self._gather()
self.loader = FileSystemLoader("templates")
self.jinja = Environment(loader=self.loader)
if environ.get("S3_IS_MINIO"):
self.upload_listing(self.create_listing(
["alpine/", "ubuntu/"]
), path_="index.html")
self.upload_listing(self.create_listing(
["3.13/", "3.12/"]
), path_="alpine/index.html")
self.upload_listing(self.create_listing(
["bionic/", "focal/"]
), path_="ubuntu/index.html")
def _gather(self) -> list:
existing_files = []
for key in self.manifest:
existing_files.append(key)
new_files = []
for file_ in listdir(self.path):
_dir = join(self.path, file_)
if isfile(_dir) and not file_.endswith("any.whl"):
file_name = basename(_dir)
if file_name not in existing_files:
new_files.append(file_name)
self.new_files = new_files
self.files = existing_files + new_files
self.files.sort(key=lambda v: v.upper())
def get_manifest(self) -> dict:
filepath = f"{self.distro}/{self.rel}/manifest.json"
try:
self.manifest = loads(self.download(self.bucket_name, filepath).decode("UTF-8"))
except botoexeptions.ClientError:
self.upload_manifest({})
self.manifest = {}
def upload_manifest(self, obj: dict):
filepath = f"{self.distro}/{self.rel}/manifest.json"
obj = bytes(dumps(obj), "ascii")
file_ = BytesIO(obj)
file_.seek(0)
self.upload(self.bucket_name, filepath, file_)
def add(self):
filepath = f"{self.distro}/{self.rel}/"
for wheel in self.new_files:
file_ = open(self.path + wheel, "rb")
self.upload(self.bucket_name, filepath + wheel, file_)
self.manifest[wheel] = {"distribution": self.distro, "release": self.rel}
self.upload_manifest(self.manifest)
self.upload_listing(self.create_listing(self.files))
def create_listing(self, wheels: list) -> str:
template = self.jinja.get_template("template.j2")
output = template.render(wheels=wheels)
return output
def upload_listing(self, obj: str, path_=None):
filepath = path_ if path_ else f"{self.distro}/{self.rel}/index.html"
obj = bytes(obj, "ascii")
file_ = BytesIO(obj)
file_.seek(0)
self.upload(self.bucket_name, filepath, file_)
if __name__ == "__main__":
parser = ArgumentParser(prog="Wheelie", description="Program to upload and present wheels to S3 buckets")
parser.add_argument("-d", "--distribution", type=str,
help="Distribution the wheels are target for.", choices=["alpine", "ubuntu"], required=True)
parser.add_argument("-r", "--release", type=str,
help="Release of the distribution.", choices=["3.12", "3.13", "focal"], required=True)
parser.add_argument("-p", "--path", type=str, default="../out/",
help="Path to the wheels")
parser.add_argument("-l", "--list", action="store_true", help="Lists files that would be added")
parser.add_argument("-a", "--add", action="store_true", help="Adds new files, and builds the index")
args = parser.parse_args()
distro = args.distribution
rel = args.release
folder = args.path
inx = Index(distribution=distro, release=rel, folder=folder)
if args.list:
print(inx.new_files if inx.new_files != [] else "No new files")
if args.add:
inx.add()