diff --git a/utils_v2/.gitignore b/utils_v2/.gitignore new file mode 100644 index 0000000..ebb370b --- /dev/null +++ b/utils_v2/.gitignore @@ -0,0 +1,8 @@ +/.venv/ +/.idea/ +**/__pycache__/ +__pycache__/ + +*.pem +*.pyc +*.pyd \ No newline at end of file diff --git a/utils_v2/__init__.py b/utils_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/__init__.py b/utils_v2/ai/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/image_classification/__init__.py b/utils_v2/ai/image_classification/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/image_classification/async_blur.py b/utils_v2/ai/image_classification/async_blur.py new file mode 100644 index 0000000..f69af42 --- /dev/null +++ b/utils_v2/ai/image_classification/async_blur.py @@ -0,0 +1,202 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Monday, 2nd Sept., 2024 + + OBJECTIVE: + + To provide an easy way to assess images for blurriness. + Tried and implemented using HuggingFace behaviour_models. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import io + +# For using the AI model: +import torch +from transformers import pipeline + +# To download images from URLs: +import requests + +# To read images: +from PIL import Image + +# Common: +from shared import variables + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AssessImageBlur: + + def __init__( + self, + model, + blur_label = "Blur", + clarity_label = "Normal", + blur_threshold = 0.25, + clarity_threshold = 0.65 + ): + + """ + This class assesses the input image and tells if it is blurry or clear. + :param model: The model to use. Either the name of the HuggingFace repo, or the directory where the model is + stored. + :param blur_threshold: The max allowed blurriness (0 to 1 range). + :param clarity_threshold: The minimum needed clarity (0 to 1 range). + """ + + # Note down the config: + self.__model = model + self.__device = "gpu" if torch.cuda.is_available() else "cpu" + self.__blur_label = blur_label + self.__clarity_label = clarity_label + self.__blur_threshold = blur_threshold + self.__clarity_threshold = clarity_threshold + + # Initialize the classifier: + self.__classifier = pipeline( + task = "image-classification", + model = self.__model, + device = self.__device + ) + + @staticmethod + async def read_image(source): + + """ + Reads an image in whichever format it is provided and returns it as a PIL object. + :param source: The image as either a path or a URL or a io.BytesIO object. + :return: The image opened as a PIL object. + """ + + # If the input image is already a PIL image: + if isinstance(source, Image.Image): return source + + # If a buffer is provided: + elif isinstance(source, io.BytesIO): + source.seek(0) + return Image.open(source) + + # If a string is provided (local path or URL): + elif isinstance(source, str): + if source.startswith("http://") or source.startswith("https://"): + response = await variables.http_client.get(source) + return Image.open(io.BytesIO(response.content)) + else: return Image.open(source) + + async def classify(self, image): + + """ + To get the prediction of the model from the given input image. + :param image: The image (as a PIL object or file path or a URL). + :return: The dictionary of classes with their respective predictions. + """ + + image = await self.read_image(image) + predictions = self.__classifier(image) + classes = {p["label"]: p["score"] for p in predictions} + return classes + + async def is_ok( + self, + image, + blur_threshold = None, + clarity_threshold = None + ): + + """ + Checks if the image given to it can be used, or should be rejected. + :param image: The image (as a PIL object or file path or a URL). + :param blur_threshold: A custom threshold to test against. If not provided, the default will be taken that was + provided when the instance was created. + :param clarity_threshold: A custom threshold to test against. If not provided, the default will be taken that + was provided when the instance was created. + :return: True if the image is okay, else False. + """ + + image = await self.read_image(image) + blur_threshold = blur_threshold or self.__blur_threshold + clarity_threshold = clarity_threshold or self.__clarity_threshold + classes = await self.classify(image) + if ( + classes[self.__blur_label] <= blur_threshold and + classes[self.__clarity_label] >= clarity_threshold + ): return True + return False + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import asyncio + + + async def main(): + image = r"/home/developer/Downloads/low-res-check.png" + assessor = AssessImageBlur(model = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/image_classification/BlurOrBokeh") + usable = await assessor.is_ok(image) + classes = await assessor.classify(image) + print("IS OKAY:", usable) + print("CLASSES:", classes) + + asyncio.run(main()) diff --git a/utils_v2/ai/image_classification/async_nsfw.py b/utils_v2/ai/image_classification/async_nsfw.py new file mode 100644 index 0000000..446c34a --- /dev/null +++ b/utils_v2/ai/image_classification/async_nsfw.py @@ -0,0 +1,188 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Monday, 2nd Sept., 2024 + + OBJECTIVE: + + To provide an easy way to assess images for adult content. + Tried and implemented using HuggingFace behaviour_models. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import io + +# For using the AI model: +import torch +from transformers import pipeline + +# To download images from URLs: +import requests + +# To read images: +from PIL import Image + +# Common: +from shared import variables + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AssessImageNSFW: + + def __init__( + self, + model, + nsfw_label = "nsfw", + nsfw_threshold = 0.25, + ): + + """ + This class assesses the input image and tells if it is blurry or clear. + :param model: The model to use. Either the name of the HuggingFace repo, or the directory where the model is + stored. + :param nsfw_threshold: The max allowed blurriness (0 to 1 range). + """ + + # Note down the config: + self.__model = model + self.__device = "gpu" if torch.cuda.is_available() else "cpu" + self.__nsfw_label = nsfw_label + self.__nsfw_threshold = nsfw_threshold + + # Initialize the classifier: + self.__classifier = pipeline( + task = "image-classification", + model = self.__model, + device = self.__device + ) + + @staticmethod + async def read_image(source): + + """ + Reads an image in whichever format it is provided and returns it as a PIL object. + :param source: The image as either a path or a URL or a io.BytesIO object. + :return: The image opened as a PIL object. + """ + + # If the input image is already a PIL image: + if isinstance(source, Image.Image): return source + + # If a buffer is provided: + elif isinstance(source, io.BytesIO): + source.seek(0) + return Image.open(source) + + # If a string is provided (local path or URL): + elif isinstance(source, str): + if source.startswith("http://") or source.startswith("https://"): + response = await variables.http_client.get(source) + return Image.open(io.BytesIO(response.content)) + else: return Image.open(source) + + async def classify(self, image): + + """ + To get the prediction of the model from the given input image. + :param image: The image (as a PIL object or file path or a URL). + :return: The dictionary of classes with their respective predictions. + """ + + image = await self.read_image(image) + predictions = self.__classifier(image) + classes = {p["label"]: p["score"] for p in predictions} + return classes + + async def is_ok( + self, + image, + nsfw_threshold = None + ): + + """ + Checks if the image given to it can be used, or should be rejected. + :param image: The image (as a PIL object or file path or a URL). + :param nsfw_threshold: A custom threshold to test against. If not provided, the default will be taken that was + provided when the instance was created. + :return: True if the image is okay, else False. + """ + + image = await self.read_image(image) + nsfw_threshold = nsfw_threshold or self.__nsfw_threshold + classes = await self.classify(image) + return True if classes[self.__nsfw_label] <= nsfw_threshold else False + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import asyncio + + + async def main(): + + assessor = AssessImageNSFW(model = r"/path/to/model") + usable = await assessor.is_ok(r"/path/to/image/img.jpg") + classes = await assessor.classify(r"https://...") + + + asyncio.run(main()) diff --git a/utils_v2/ai/image_segmentation/__init__.py b/utils_v2/ai/image_segmentation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/image_segmentation/async_dichotomous_segmentation.py b/utils_v2/ai/image_segmentation/async_dichotomous_segmentation.py new file mode 100644 index 0000000..fb6d6c7 --- /dev/null +++ b/utils_v2/ai/image_segmentation/async_dichotomous_segmentation.py @@ -0,0 +1,223 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Saturday, 14th Sept., 2024 + + OBJECTIVE: + + To provide an easy way to get masks from dichotomous image segmentation. + + This code uses a very specific model from HuggingFace: "ZhengPeng7/BiRefNet-portrait". + You may experiment with other behaviour_models too, but make sure that model is made for "dichotomous" behaviour. This + means that the model should have only two classes like "foreground", and "background". The specified model was + trained for implementing portrait mode style blurring of backgrounds. + + The originally tested model has an MIT license as per their GitHub page. The code in this file may or may not + support drop-in replacement for other behaviour_models, please be aware about this. + + REFERENCES: + + 01. https://huggingface.co/ZhengPeng7/BiRefNet-portrait + 02. https://github.com/ZhengPeng7/BiRefNet + + DOWNLOADS: + + 01. https://huggingface.co/ZhengPeng7/BiRefNet-portrait/tree/main + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# For system-level activities: +import io + +# To work with PIL images: +from PIL import Image + +# To make asynchronous HTTP calls: +from shared.variables import http_client + +# For using the AI model: +import torch +from torchvision import transforms +from transformers import AutoModelForImageSegmentation + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class DichotomousSegmenter: + + def __init__( + self, + model = r"zhengpeng7/BiRefNet-portrait" + ): + + """ + This class uses Dichotomous Image Segmentation to produce a mask of what is in the foreground (or what is of + interest in a given image). + :param model: The path of the model. + """ + + # Check if GPU is available for faster predictions: + self.__cuda_is_available = True if torch.cuda.is_available() else False + + # Initialize the model: + self.__model = AutoModelForImageSegmentation.from_pretrained( + pretrained_model_name_or_path = model, + trust_remote_code = True + ) + torch.set_float32_matmul_precision(["high", "highest"][0]) + self.__model.to("cuda" if self.__cuda_is_available else "cpu") + + @staticmethod + async def read_image(source): + + """ + Reads an image in whichever format it is provided and returns it as a PIL object. + :param source: The image as either a path or a URL or a io.BytesIO object. + :return: The image opened as a PIL object. + """ + + # If the input image is already a PIL image: + if isinstance(source, Image.Image): return source + + # If a buffer is provided: + elif isinstance(source, io.BytesIO): + source.seek(0) + return Image.open(source) + + # If a string is provided (local path or URL): + elif isinstance(source, str): + if source.startswith("http://") or source.startswith("https://"): + response = await http_client.get(source) + return Image.open(io.BytesIO(response.content)) + else: return Image.open(source) + + async def get_mask(self, image): + + """ + Process the image to figure out the mask. In some cases (like the image of a sunset) you will have no white + pixels in the mask. White pixels represent the areas that the AI considered to be the foreground (area of + interest) and the black parts are the background. + :param image: The image data as either a PIL object, or a path to a file on the local disk, or a URL. + :return: The mask as a PIL object. + """ + + # Prepare the image transformer: + image_size = (1024, 1024) + transform_image = transforms.Compose([ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ]) + + # Read and transform the image: + image = await self.read_image(image) + input_image = transform_image(image).unsqueeze(0).to("cuda" if self.__cuda_is_available else "cpu") + + # Make the prediction: + with torch.no_grad(): + if self.__cuda_is_available: mask = self.__model(input_image)[-1].sigmoid().cuda() + else: mask = self.__model(input_image)[-1].sigmoid().cpu() + + # Create a PIL object of the mask data: + mask = mask[0].squeeze() + mask = transforms.ToPILImage()(mask) + mask = mask.resize(image.size) + + # Done here: + return mask + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import asyncio + import numpy as np + + async def main(): + + model = DichotomousSegmenter() + + images = [ + # r"/home/developer/Downloads/kate.jpg", + # r"/home/developer/Downloads/IMG-20240730-WA0001.jpg", + # r"/home/developer/Downloads/pexels-pixabay-57416.jpg", + # r"/home/developer/Downloads/summer_clothes.jpg", + # r"/home/developer/Downloads/low_res_cat - upscaled.png", + # r"/home/developer/Downloads/blurry_traffic.jpg", + r"/home/developer/Downloads/card_0.jpg", + r"/home/developer/Downloads/card_1.jpg", + r"/home/developer/Downloads/card_2.jpg", + r"/home/developer/Downloads/card_3.jpg", + r"/home/developer/Downloads/sushmita_card.jpg", + r"/home/developer/Downloads/niranjan_card.jpg", + r"/home/developer/Downloads/niranjan_card_2.jpg", + ] + + for image in images: + print(image.split("/")[-1]) + image = await DichotomousSegmenter.read_image(image) + mask = await model.get_mask(image) + image.show() + mask.show() + print("MASK!") + array = np.array(mask) + print("Dimensions of the array:", array.shape) + masked_count = np.sum(array > 125) + print(masked_count, type(masked_count), np.size, type(np.size)) + percentage = masked_count / np.prod(array.shape) + print("Area:", round(percentage * 100, 4)) + # image.putalpha(mask) + # image.show() + # mask.show() + + + asyncio.run(main()) diff --git a/utils_v2/ai/object_detection/__init__.py b/utils_v2/ai/object_detection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/object_detection/yolo.py b/utils_v2/ai/object_detection/yolo.py new file mode 100644 index 0000000..a435ca6 --- /dev/null +++ b/utils_v2/ai/object_detection/yolo.py @@ -0,0 +1,176 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 4th Jul, 2024 + + OBJECTIVE: + + To provide a class to detect objects in images using YOLO behaviour_models. + + REFERENCES: + + 1) Code Examples: https://docs.ultralytics.com/usage/python/ + 2) Models: https://docs.ultralytics.com/models/yolov8/#supported-tasks-and-modes + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import os + +# To use YOLO architecture: +from ultralytics import YOLO + +# To work with images: +from PIL import Image + +# utils: +from utils_v2.string import json + +# Common: +from shared import constants + +# For debugging: +from icecream import IceCreamDebugger + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +class YoloDetect: + + def __init__( + self, + model_file, + debug = True, + debug_prefix = "YOLO | " + ): + + # Load the model: + self._model = YOLO(model = model_file) + + # Initialize the dbugging tool: + self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) + if not debug: self._printer.disable() + + def disable_debug(self): + self._printer.disable() + + def enable_debug(self): + self._printer.enable() + + @staticmethod + def _process_result(raw_result): + + # Extract the various types of results available from the inference: + speed = {k: (v / 100.0) for k, v in raw_result.speed.items()} + class_mapping = raw_result.names + probabilities = raw_result.probs + boxes = raw_result.boxes + masks = raw_result.masks + + # Construct a default response: + response_json = { + "speed": speed, + "classes": class_mapping, + "boxes": None + } + + # Object-Detection results: + formatted_boxes = [] + for box in boxes: + detected_class = int(box.cls.numpy()[0]) + x1, y1, x2, y2 = box.xyxy[0] + formatted_boxes.append( + { + "class": detected_class, + "className": class_mapping[detected_class], + "confidence": float(box.conf.numpy()[0]), + "x1": int(x1.numpy()), + "y1": int(y1.numpy()), + "x2": int(x2.numpy()), + "y2": int(y2.numpy()) + } + ) + response_json["boxes"] = formatted_boxes + + # Done here: + return response_json + + def predict(self, image, show = False, verbose = False): + + # Process the input image with the given task: + results = self._model( + source = image, + show = show, + verbose = verbose + ) + + # Based on the task, interpret the results: + results_json = self._process_result(results[0]) + + # Done here: + return results_json + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + # model_file_path = os.path.join(constants.PROJECT_DIRECTORY, "ai", "yolo", "behaviour_models", "yolov8x.pt") + model_file_path = r"/home/developer/PycharmProjects/utils/data/ai/models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt" + # sample_image_path = r"/home/developer/Downloads/2_cats.jpg" + sample_image_path = r"/home/developer/Downloads/flattened_image.jpg" + my_yolo = YoloDetect(model_file = model_file_path) + results = my_yolo.predict(image = Image.open(sample_image_path)) + print("FINAL RESULTS:", json.to_string(results)) + diff --git a/utils_v2/ai/tts/__init__.py b/utils_v2/ai/tts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/ai/tts/easy_tts.py b/utils_v2/ai/tts/easy_tts.py new file mode 100644 index 0000000..fc4bb9b --- /dev/null +++ b/utils_v2/ai/tts/easy_tts.py @@ -0,0 +1,175 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 26th Sept., 2024 + + OBJECTIVE: + + To provide a fast way to make audio from TTS engines. + + REFERENCES: + + 01. Usage: https://github.com/myshell-ai/MeloTTS/blob/main/docs/install.md#python-api + 02. Installation: https://github.com/myshell-ai/MeloTTS/blob/main/docs/install.md#linux-and-macos-install + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import io + +# To use the AI: +from melo.api import TTS +import numpy as np +import soundfile + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class EasyTTS: + + __sampling_rate = 44_100 + + def __init__( + self, + language = "EN", + speed = 1.0 + ): + + self.__audio = np.zeros(1) + self.__language = language + self.__speed = speed + self.__model = TTS(language = language, device = "auto") + self.__speakers = self.__model.hps.data.spk2id + + def list_speakers(self): + return list(self.__speakers.keys()) + + def speak( + self, + text, + speaker + ): + + this_audio = self.__model.tts_to_file( + text, + self.__speakers[speaker], + speed = self.__speed, + quiet = True + ) + + self.__audio = np.concatenate((self.__audio, this_audio)) + + def pause(self, seconds): + + self.__audio = np.concatenate(( + self.__audio, + np.zeros(int(self.__sampling_rate * seconds)) + )) + + def to_wav(self, path = None): + + if path is None: + audio_buffer = io.BytesIO() + soundfile.write(audio_buffer, self.__audio, self.__sampling_rate) + return audio_buffer + + else: soundfile.write(path, self.__audio, self.__sampling_rate) + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + speaker = "EN-BR" + tts_maker = EasyTTS(language = "EN", speed = 0.9) + + tts_maker.speak( + text = """ + Imagine delighting your doctors with a personalized calendar, + crafted from their own cherished memories. + """, + speaker = speaker + ) + tts_maker.pause(seconds = 0.3) + tts_maker.speak( + text = """ + Every day, as they turn the page, + they’ll not only relive those special moments but also remember you, + the one who made it happen. + """, + speaker = speaker + ) + tts_maker.pause(seconds = 0.75) + tts_maker.speak( + text = "STEP 1:", + speaker = speaker + ) + tts_maker.pause(seconds = 0.3) + tts_maker.speak( + text = "Start by identifying the doctors you’d like to engage with, and add them to our app.", + speaker = speaker + ) + tts_maker.pause(seconds = 0.3) + tts_maker.speak( + text = "No rush, you can add their photographs later as well.", + speaker = speaker + ) + tts_maker.pause(seconds = 0.3) + tts_maker.speak( + text = "With this, your engagement funnel is created.", + speaker = speaker + ) + + tts_maker.to_wav(r"/home/developer/Downloads/voiceover.wav") diff --git a/utils_v2/ai/tts/test.wav b/utils_v2/ai/tts/test.wav new file mode 100644 index 0000000..d38cfc0 Binary files /dev/null and b/utils_v2/ai/tts/test.wav differ diff --git a/utils_v2/api/__init__.py b/utils_v2/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/api/__pycache__/__init__.cpython-310.pyc b/utils_v2/api/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..c5c8546 Binary files /dev/null and b/utils_v2/api/__pycache__/__init__.cpython-310.pyc differ diff --git a/utils_v2/api/__pycache__/async_quart.cpython-310.pyc b/utils_v2/api/__pycache__/async_quart.cpython-310.pyc new file mode 100644 index 0000000..74b15d2 Binary files /dev/null and b/utils_v2/api/__pycache__/async_quart.cpython-310.pyc differ diff --git a/utils_v2/api/__pycache__/codes.cpython-310.pyc b/utils_v2/api/__pycache__/codes.cpython-310.pyc new file mode 100644 index 0000000..2977ca4 Binary files /dev/null and b/utils_v2/api/__pycache__/codes.cpython-310.pyc differ diff --git a/utils_v2/api/__pycache__/log.cpython-310.pyc b/utils_v2/api/__pycache__/log.cpython-310.pyc new file mode 100644 index 0000000..38cf927 Binary files /dev/null and b/utils_v2/api/__pycache__/log.cpython-310.pyc differ diff --git a/utils_v2/api/__pycache__/metrics_prometheus.cpython-310.pyc b/utils_v2/api/__pycache__/metrics_prometheus.cpython-310.pyc new file mode 100644 index 0000000..724908a Binary files /dev/null and b/utils_v2/api/__pycache__/metrics_prometheus.cpython-310.pyc differ diff --git a/utils_v2/api/__pycache__/response.cpython-310.pyc b/utils_v2/api/__pycache__/response.cpython-310.pyc new file mode 100644 index 0000000..dd15b88 Binary files /dev/null and b/utils_v2/api/__pycache__/response.cpython-310.pyc differ diff --git a/utils_v2/api/async_quart.py b/utils_v2/api/async_quart.py new file mode 100644 index 0000000..2158d7e --- /dev/null +++ b/utils_v2/api/async_quart.py @@ -0,0 +1,1264 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Friday, 2nd Aug., 2024 + + OBJECTIVE: + + To provide an easy way to perform repetitive tasks in quart. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + + NOTES: + + 01. PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING ANY OF THESE + DECORATORS. + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To work with Quart: +from quart import request, current_app, g, make_response + +# To make decorators: +from functools import wraps + +# My utils: +from utils_v2.string import json +from utils_v2.datetime import datetime +from utils_v2.security import sanitizers +from utils_v2.api.codes import StatusCodes, HttpCodes +from utils_v2.api.log import APILogModel +from utils_v2.api.response import ResponseModel +from utils_v2.logging.context import AsyncLoggerContext, AsyncMongoLogger + +# System-level activities: +import io +import os + +# To work with datatypes: +from types import NoneType +import pandas as pd + +# For Pydantic data-behaviour_models: +import pydantic +from typing import List + +# For hashing and shortening the hash: +import hashlib +import base64 + +# To make things human-readable: +import humanize + +# For debugging: +import traceback +import random +import string + +# For asynchronous activities: +import asyncio + +# For timekeeping: +import time + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# Chars to choose from for random strings: +ALPHANUMERIC_CHARS = string.ascii_letters + string.digits + +# To capture system information: +PROCESS_ID = os.getppid() +PARENT_PROCESS_ID = os.getppid() + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** EXCEPTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +class AuthDetailsIncompleteException(Exception): + def __str__(self): + return "incomplete auth details" + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +async def make_ordered_json(json_data, http_code = 200): + + """ + Quart sorts the fields of a dict when converting to a JSON response. Here we are manually making the response when + the sequence of the fields is sensitive. + :param json_data: The data (dict, list, etc.) to be converted to a JSON string. + :param http_code: The HTTP status code you want to send with the response. + :return: The JSON-ified response such that the sequence of the fields is maintained. + """ + + response = await make_response(json.to_string(json_data, no_space = True), http_code) + response.headers["Content-Type"] = "application/json" + return response + + +# --------------------------------------------------------------------------------------------------------------------- + + +async def data_from_request(inbound_request): + + """ + Adaptively extract the params from incoming request in whichever way it was provided. + :param inbound_request: The request that came in. + :return: A dict (could be empty) of the data/params that came in with the request. + """ + + inbound_data = {} + + # Extract data from the params in the URL: + from_args = inbound_request.args.to_dict() + if isinstance(from_args, dict): + for k, v in from_args.items(): + inbound_data[k] = v + + # Extract data from the raw JSON data: + from_json = await request.get_json() + if isinstance(from_json, dict): + for k, v in from_json.items(): + inbound_data[k] = v + + # Extract inputs from the form body: + from_form = await inbound_request.form + from_form = from_form.to_dict() + if isinstance(from_form, dict): + for k, v in from_form.items(): + inbound_data[k] = v + + return inbound_data + + +# --------------------------------------------------------------------------------------------------------------------- + + +async def file_from_request(inbound_request, file_key): + + """ + Extracts ONE file from the incoming request's form-data. + WARNING: NOT RECOMMENDED FOR LARGE FILES. STRICTLY USE FOR SMALL FILES THAT WON'T CRASH THE SCRIPT. + :param inbound_request: The incoming request. + :param file_key: The key of the file that you want to extract. + :return: A tuple of the file's name and data. + """ + + file_name = None + file_data = None + files = await inbound_request.files + + if file_key in files: + file_name = files[file_key].filename + file_data = io.BytesIO(files[file_key].read()) + + return file_name, file_data + + +# --------------------------------------------------------------------------------------------------------------------- + + +async def files_from_request(inbound_request: request): + + """ + Extracts ALL files from the incoming request's form-data. + WARNING: NOT RECOMMENDED FOR LARGE FILES. STRICTLY USE FOR SMALL FILES THAT WON'T CRASH THE SCRIPT. + :param inbound_request: The incoming request. + :return: A dict describing the file's name, data, and size. + """ + + files = await inbound_request.files + + inbound_files = {} + for file_key in files: + file_data = io.BytesIO(files[file_key].read()) + file_size = file_data.seek(0, 2) + file_data.seek(0) + inbound_files[file_key] = { + "name": files[file_key].filename, + "data": file_data, + "size": file_size, + "type": files[file_key].content_type + } + + return inbound_files + + +# --------------------------------------------------------------------------------------------------------------------- + + +async def headers_from_request( + inbound_request, + mandatory_keys: list = None +): + + """ + Extract custom headers and some extra info. from the incoming request. + Raises an exception if any mandatory key is missing. + IMPORTANT: CUSTOMIZE THIS FOR THE NEEDS OF YOUR PROJECT. + :param inbound_request: The request that came in. + :param mandatory_keys: The keys that you need to have in the auth. + :return: The extracted auth details. + """ + + # Start by extracting whatever complies with the format of "X-{Header-Name}": + head_json = {k: v for k, v in request.headers.items() if k.startswith("X-")} + + # Now note down things that you want to keep from incoming requests: + head_json["Remote-IP"] = inbound_request.remote_addr + head_json["Host"] = inbound_request.headers.get("Host") + head_json["Origin"] = inbound_request.headers.get("Origin") + head_json["User-Agent"] = inbound_request.headers.get("User-Agent") + head_json["Content-Type"] = inbound_request.headers.get("Content-Type") + head_json["Content-Length"] = inbound_request.headers.get("Content-Length") + + # Raise an exception if any of the mandatory auth details were missing: + if mandatory_keys is not None: + available_keys = head_json.keys() + for mandatory_key in mandatory_keys: + if mandatory_key not in available_keys: raise AuthDetailsIncompleteException + + # Done here: + return head_json + + +# --------------------------------------------------------------------------------------------------------------------- + + +def cause_exception(): + + """ + Call this from any function when you want to raise an exception. + Example use case would be when receiving data from an API call and that field is not supposed to be null. + :return: None. + """ + + return 100/0 + + +# --------------------------------------------------------------------------------------------------------------------- + + +def describe_exception(exc): + + """ + Describes the exception in detail. It extracts the type of exception, a brief message, and even the entire + traceback. Useful for debugging in details without the terminal. You could either log the resultant dict or send it + to the dev team over some service like WhatsApp/Telegram. + :param exc: The exception that occurred. + :return: The dict that explains the exception. + """ + + exc_desc = { + "type": type(exc).__name__, + "msg": str(exc), + "tb": [str(exc_tb) for exc_tb in traceback.format_exception(exc, value = exc, tb = exc.__traceback__)] + } + + return exc_desc + + +# --------------------------------------------------------------------------------------------------------------------- + + +def messages_from_pydantic_exception(exception, as_str = True, sep = ", "): + + """ + Creates a list of readable error messages from Pydantic's validation failure. + :param exception: Pydantic's ValidationError + :param as_str: Set to True to receive all messages as one string, False to receive an array of strings. + :param sep: The separator to use when joining multiple messages as one string. + :return: A list of messages of all the things that went wrong. + """ + + # Make a variable to hold all individual messages: + messages = [] + + # Interpret all the problems: + for error in exception.errors(): + loc = " --> ".join([str(item) for item in error["loc"]]) + msg = error.get("msg") + if error["type"] == "missing": messages.append(f"missing input: '{loc}'") + elif error["type"] == "model_type": messages.append(f"invalid input: '{loc}'") + elif error["type"] == "bool_parsing": messages.append(f"invalid bool: '{loc}'") + elif error["type"] == "string_type": messages.append(f"invalid string: '{loc}'") + elif error["type"] == "float_parsing": messages.append(f"invalid float: '{loc}'") + elif error["type"] == "int_parsing": messages.append(f"invalid integer: '{loc}'") + elif error["type"] == "extra_forbidden": messages.append(f"extra input: '{loc}'") + elif error["type"] == "value_error": messages.append(f"validation failed: '{loc}' ({msg})") + else: messages.append(f"invalid value: '{loc}'") + + # Return a response as per the preference of the user: + if as_str: return sep.join(messages) + else: return messages + + +# --------------------------------------------------------------------------------------------------------------------- + + +def summarize_variable( + value, + str_limit = 100, + expand: bool | int = False, + sensitive_keys: list[str] = None +): + + """ + To summarize an input value to capture the essence without hoarding to much data. + :param value: Anything that you want to summarize. + :param str_limit: The max. no. of chars of a string to retain. + :param expand: Set to True for full expansion, False for no expansion, and an integer for a specific level of + expansion. Applicable on iterables and dicts. The smaller this number, the more concise the summary will be, + and vice versa. + :param sensitive_keys: The list of keys (of a dict) to obscure when summarizing. + :return: The summarized version of the input. + """ + + # If a null value was sent: + if value is None: return None + + # Check the sensitive keys: + if sensitive_keys is None: sensitive_keys = [] + + # Handle datatypes that you don't want to modify: + if isinstance(value, (int, float, bool, NoneType)): pass + + # When the value is a list or similar iterable: + elif isinstance(value, (list, tuple, set)): + if expand: + if not isinstance(expand, bool): expand -= 1 + value = [AsyncLoggerContext.summarize( + v, + expand = expand, + sensitive_keys = sensitive_keys + ) for v in value] + else: value = f"array of {len(value)} item(s)" + + # If the value is a dict: + elif isinstance(value, dict): + if expand: + if not isinstance(expand, bool): expand -= 1 + value = { + k: AsyncLoggerContext.summarize( + v, + expand = expand, + sensitive_keys = sensitive_keys + ) if k not in sensitive_keys else "********" + for k, v in value.items() + } + else: value = f"object of {len(value.keys())} field(s) [{', '.join(value.keys())}]" + + # When a dataframe is passed: + elif isinstance(value, pd.DataFrame): + cols = value.columns.to_list() + value = f"table with {len(cols)} col(s) [{', '.join(cols)}] and {len(value)} row(s)" + str_limit = 999 + + # If the input is some form of non-standard object: + else: value = str(value) + + # Handle strings: + if isinstance(value, str): + if len(value) > str_limit: value = value[:str_limit] + "..." + + # Done here: + return value + + +# --------------------------------------------------------------------------------------------------------------------- + + +def set_api_version(api_version): + + """ + Use this decorator to automatically note down the API version no. and propagate it throughout the downstream + decorators. Use this as the entry point if possible. + :param api_version: The version code to assign to the API. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # set the version information in the variable. + # This makes it available to the downstream decorators too! + kwargs["api_version"] = api_version + + # we are ready to call the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Add the API version to the response: + if isinstance(response, ResponseModel): response.api_version = api_version + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def read_input( + read_headers = True, + read_data = True, + read_files = True, + sanitize_headers = True, + sanitize_data = True +): + + """ + Use this decorator to read the inputs from the incoming request and sanitize them. Sanitization makes the inputs + safe against certain threats like injections attacks. If you expect to take in inputs that you want to use to run + database commands, you could disable them manually. + PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. + :param read_headers: Whether, or not, to read the headers of the request. + :param read_data: Whether, or not, to read the data of the request. + :param read_files: Whether, or not, to read the files of the request. + :param sanitize_headers: Whether, or not, you would like to sanitize the params coming in through the headers. + :param sanitize_data: Whether, or not, you would like to sanitize the params coming in through the body or query. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Get the headers: + if read_headers: + kwargs["inbound_headers"] = await headers_from_request(request) + if sanitize_headers: kwargs["inbound_headers"] = sanitizers.for_mongo(kwargs["inbound_headers"]) + # else: kwargs["inbound_headers"] = None + + # Get the data: + if read_data: + kwargs["inbound_data"] = await data_from_request(request) + if sanitize_data: kwargs["inbound_data"] = sanitizers.for_mongo(kwargs["inbound_data"]) + # else: kwargs["inbound_data"] = None + + # Get small files from the request: + if read_files: kwargs["inbound_files"] = await files_from_request(request) + # else: kwargs["inbound_files"] = None + + # # We also make a provision for capturing an identifier + # # for the logs that we make through a sister decorator: + # kwargs["log_id"] = "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(8)) + # kwargs["log_chain"] = kwargs.get("inbound_headers", {}).get("X-Log-Chain") or kwargs["log_id"] + + # Now that we have unpacked the incoming data, + # we are ready to run the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def get_session_info( + key: str, + model: str = None, + session_func: str = None, + session_coro: str = None, + get: str = None, + mandatory: bool = False, + sensitive_keys: List[str] = None +): + + """ + To get the information about the user from the session token or some similar identifier. Use this after + 'read_input', and note that the callable func/coro should only take on parameter - the thing to identify the session + by, and it should return a dictionary with all the needed details. + PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. + :param key: The key in either the 'inbound_data' or the 'inbound_headers' from which the session info will be + available. Examples: 'sessionToken' or 'X-Session-Token'. + :param model: The object from which the func/coro should be called. This could be the name of the variable holding + an instance of a class. Should be available in the scope of 'current_app'. + :param session_func: The synchronous func of the model to call. This can either be an independent function or a + class's method. Should be available in the scope of 'current_app'. + :param session_coro: The asynchronous func of the model to call. This can either be an independent function or a + class's method. Ignored if 'func' was provided. Should be available in the scope of 'current_app'. + :param get: The field inside the dict to get. Specify this in dot notation. The whole response will be returned as + is if this is null. + :param mandatory: If set to True, the API call will enforce session checking; and if a session is not found, the + client will receive an unauthorized failure message. + :param sensitive_keys: Any keys to obscure when retrieving details about the session. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Start by assuming failure: + session_info = None + + # Get the identifier of the session from the values extracted in the 'read_input' decorator: + session_id = kwargs["inbound_data"].get(key, kwargs["inbound_headers"].get(key)) + + # If the key is present, + # we call the function that will get the session's info for us: + if session_id is not None: + + try: + + # For synchronous functions: + if session_func: + if model: session_info = getattr(getattr(current_app, model), session_func)(session_id) + else: session_info = getattr(current_app, session_func)(session_id) + + # For asynchronous coroutines: + elif session_coro: + if model: session_info = await getattr(getattr(current_app, model), session_coro)(session_id) + else: session_info = await getattr(current_app, session_coro)(session_id) + + # If something goes wrong: + except Exception as exception: + if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) + + # We note down whatever we got: + if session_info and isinstance(get, str): + for subkey in get.split("."): + session_info = session_info.get(subkey, {}) if isinstance(session_info, dict) else {} + session_info = summarize_variable(session_info, expand = True, sensitive_keys = sensitive_keys) + kwargs["session_info"] = session_info + + # If no session info was found, but it was mandatory: + if mandatory and not session_info: + response = ResponseModel( + status_code = StatusCodes.FAILED, + http_code = HttpCodes.UNAUTHORIZED, + message = f"invalid session" + ) + + # Now we are ready to run the function that we are wrapping: + else: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def validate_input( + mandatory_header_keys = None, + mandatory_data_keys = None, + mandatory_file_keys = None, + header_validator = None, + data_validator = None +): + + """ + USE THIS ONLY AFTER YOU HAVE USED 'read_input'. This decorator will help you run validation on the inputs that were + extracted from the request. The mandatory keys will be checked first and the validations will be run after that. If + your validator already checks for keys, you may skip mentioning mandatory keys. DO NOTE THAT YOUR VALIDATOR + FUNCTIONS MUST RAISE AN EXCEPTION FOR THIS DECORATOR TO WORK. + :param mandatory_header_keys: The keys in 'inbound_headers' that are absolutely necessary. + :param mandatory_data_keys: The keys in 'inbound_data' that are absolutely necessary. + :param mandatory_file_keys: The keys in 'inbound_files' that are absolutely necessary. + :param header_validator: The function to use to validate the 'inbound_headers'. + :param data_validator: The function to use to validate 'inbound_data'. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # We first validate the mandatory header keys. + # Having a null value in this case is NOT allowed: + if mandatory_header_keys is not None: + for mandatory_key in mandatory_header_keys: + if kwargs["inbound_headers"].get(mandatory_key) is None: + return ResponseModel( + status_code = StatusCodes.HEADERS_INCOMPLETE, + message = f"missing: '{mandatory_key}'" + ) + + # Return with failure if any of the mandatory JSON details are missing. + # Having a null value is allowed, but is should be sent by the user on intention. + if mandatory_data_keys is not None: + for mandatory_key in mandatory_data_keys: + try: kwargs["inbound_data"][mandatory_key] + except: return ResponseModel( + status_code = StatusCodes.DATA_INCOMPLETE, + message = f"missing: '{mandatory_key}'" + ) + + # Return with failure if any of the mandatory file-keys details are missing: + if mandatory_file_keys is not None: + provided_file_keys = kwargs["inbound_files"].keys() + for mandatory_key in mandatory_file_keys: + if mandatory_key not in provided_file_keys: + return ResponseModel( + status_code = StatusCodes.FILE_MISSING, + message = f"missing: '{mandatory_key}'" + ) + + # Next we validate the headers: + if header_validator is not None: + + # Try validate the data: + try: kwargs["inbound_headers"] = header_validator(kwargs["inbound_headers"]) + + # In case some needed field is missing: + except KeyError as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = "missing: " + str(exception), + http_code = HttpCodes.BAD_REQUEST, + ) + + # In case some pydantic data model fails validation: + except pydantic.ValidationError as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = messages_from_pydantic_exception(exception), + http_code = HttpCodes.BAD_REQUEST + ) + + # In case some other exception was raised: + except Exception as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = str(exception), + http_code = HttpCodes.BAD_REQUEST + ) + + # Finally, we validate the incoming data: + if data_validator is not None: + + # Try validate the data: + try: kwargs["inbound_data"] = data_validator(kwargs["inbound_data"]) + + # In case some needed field is missing: + except KeyError as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = "missing: " + str(exception), + http_code = HttpCodes.BAD_REQUEST + ) + + # In case some pydantic data model fails validation: + except pydantic.ValidationError as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = messages_from_pydantic_exception(exception), + http_code = HttpCodes.BAD_REQUEST + ) + + # In case some other exception was raised: + except Exception as exception: + return ResponseModel( + status_code = StatusCodes.DATA_VALIDATION_FAILURE, + message = str(exception), + http_code = HttpCodes.BAD_REQUEST + ) + + # Now that we have unpacked the incoming data, + # we are ready to run the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def log_request_to_mongo( + attr_name, + collection: str = "logs", + api_version: str = None, + project: str = None, + log_type: str = None, + operation: str = None, + log_input: bool | int = True, + log_output: bool | int = True, + sensitive_keys: list = None +): + + """ + USE THIS ONLY AFTER YOU HAVE USED 'read_input'. This decorator will log the whole process of the API call to + MongoDB. The variable that holds the instance of 'AsyncMongo' needs to be accessible in the scope of 'current_app'. + PLEASE USE "ResponseModel" AS THE RETURNED VALUE OF THE API ENDPOINT IF YOU ARE USING THIS DECORATOR. + :param attr_name: The name of the variable that holds the instance of 'AsyncMongo'. It should be accessible in the + scope of 'current_app'. + :param collection: The name of the collection to write the log into. + :param api_version: The version code of the API endpoint that is being logged. + :param project: The name of the project that the endpoint was built for. + :param log_type: A hint to identify what the log was for. + :param operation: A hint to identify what was action was being performed. + :param log_input: Set to True to capture everything that went into the function, False to capture the least + info, and set it to an integer to capture a certain depth of the input (applicable on iterables and dicts. + :param log_output: The same as 'log_input', but applicable to the response from the function. + :param sensitive_keys: The list of keys to not log. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Create ids for getting logs: + kwargs["log_id"] = "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(8)) + kwargs["log_chain"] = kwargs.get("inbound_headers", {}).get("X-Log-Chain") or kwargs["log_id"] + + # Set the api version as needed: + kwargs["api_version"] = kwargs.get("api_version", api_version) + + # Make variables and extract available info.: + exception = None + response = None + request_ts = datetime.get_current_utc_date_time() + start_ts = time.perf_counter() + cpu_start_ts = time.process_time() + + # Execute the function that is being wrapped: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + try: response = await func(*args, **kwargs) + except Exception as exc: exception = exc + kwargs["decorator_count"] -= 1 + + # Ensure that the response is not null: + response = response if response is not None else ResponseModel( + status_code = StatusCodes.UNKNOWN_ERROR, + message = "null response for request" + ) + + # Add params to the response. + # THIS IS ONLY APPLICABLE WHEN THE TYPE OF THE RESPONSE IS 'ResponseModel': + if isinstance(response, ResponseModel): + response.api_version = kwargs.get("api_version") + response.log_id = kwargs.get("log_id") + + # Extract the response to log: + response_to_log = "not logged" + http_code_to_log = 200 + if log_output: + if isinstance(response, ResponseModel): response_to_log, http_code_to_log = response.for_quart() + elif isinstance(response, tuple): response_to_log, http_code_to_log = response + else: response_to_log, http_code_to_log = str(response), 200 + response_to_log = summarize_variable( + response_to_log, + expand = log_output, + sensitive_keys = sensitive_keys + ) + + # Try to get the information about the request. + # There will be no data in any of these if the decorator was used to catch start-up and shut-down events. + request_method = None + request_url = None + request_route = None + try: + request_method = f"{request.method}" + request_url = f"{request.url}" + request_route = str(request.url_rule.rule) + except: pass + + # Redact the sensitive keys: + if sensitive_keys: + for k in sensitive_keys: + for var in ["inbound_headers", "inbound_data"]: + try: kwargs[var][k] = len(str(kwargs[var][k])) * "*" + except: pass + + # Construct the log: + api_log = APILogModel( + pid = PROCESS_ID, + ppid = PARENT_PROCESS_ID, + project = project, + log = log_type, + operation = operation, + apiVer = kwargs.get("api_version"), + logId = kwargs.get("log_id"), + logChain = kwargs.get("log_chain", kwargs.get("inbound_headers", {}).get("X-Log-Chain")), + ts = request_ts, + tat = time.perf_counter() - start_ts, + cpuTime = time.process_time() - cpu_start_ts, + sessionInfo = kwargs.get("session_info"), + method = request_method, + url = request_url, + route = request_route, + headers = kwargs.get("inbound_headers"), + data = summarize_variable( + kwargs.get("inbound_data"), + expand = log_input, + sensitive_keys = sensitive_keys + ), + files = { + k: { + "name": v["name"], + "size": v["size"] + } for k, v in kwargs.get("inbound_files", {}).items() + }, + exception = None if exception is None else describe_exception(exception), + response = response_to_log, + httpCode = http_code_to_log + ) + + # Write the log: + app_attr = getattr(current_app, attr_name) + inserted_id = await app_attr.insert_one( + collection = collection, + document = api_log.model_dump() + ) + + # Return the response from the wrapped function. + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def log_chain_to_mongo(attr_name): + + """ + Use this to run everything in the context of the logging decorator. Everything that is decorated with the + custom decorator will be logged if this is used. + :param attr_name: The name of the db connection (an instance of 'AsyncMongo' or 'AsyncMongoStorage') + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Let the next in-line decorator know that it has been wrapped: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + + # Fetch the attribute and get it to log the whole chain: + async with AsyncLoggerContext.logging_context( + logger = AsyncMongoLogger(getattr(current_app, attr_name)), + log_chain = kwargs.get("log_chain") + ): + response = await func(*args, **kwargs) + + # Done here: + kwargs["decorator_count"] -= 1 + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def should_not_be_under_maintenance(attr_name): + + """ + Use this decorator to reject a request when the app is being marked as "under-maintenance". You will need to create + a boolean variable within the scope of the 'current_app' for this to work. An alternate to this is to set the value + in an environment variable named 'IS_UNDER_MAINTENANCE' to a string value of either 'True' or 'False' for + multi-worker deployments. + :param attr_name: The name of the boolean variable that will hold the information about the app being under + maintenance. If its value is True at the time of checking, the incoming request will be rejected. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Get the attribute and check if it indicates that the app is under maintenance, + # call the wrapped function if not under maintenance: + app_attr = getattr(current_app, attr_name) + env_attr = True if os.environ.get("IS_UNDER_MAINTENANCE", "False").lower() == "true" else False + if app_attr or env_attr: + response = ResponseModel(status_code = StatusCodes.DOWN_FOR_MAINTENANCE).for_quart() + else: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def only_whitelisted_ips(attr_name): + + """ + Use this decorator to reject any requests coming from unauthorized IPs. The list of IP addresses to allow must be + in a list that is accessible in the context of 'current_app'. + :param attr_name: The name of the boolean variable that will hold the information about the app being under + maintenance. If its value is True at the time of checking, the incoming request will be rejected. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Get the attribute and check if the request's IP is in the permitted list: + app_attr = getattr(current_app, attr_name) + if request.remote_addr not in app_attr: + return ResponseModel( + status_code = StatusCodes.AUTHORIZATION_FAILED, + message = "bad ip", + http_code = HttpCodes.UNAUTHORIZED + ).for_quart() + + # Now that we have checked that the IP is permitted, + # we are ready to run the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def limit_rate( + attr_name, + rate_limit: int = 5, + seconds: float = 1.0, + message = None, + header_keys: list = None, + data_keys: list = None, + allow_if_exception = False, + count_for_http_codes = None +): + + """ + Use this decorator to apply rate-limiting to incoming requests. + SUGGESTION: WHEN STACKING UP MANY RATE LIMITS, PUT THE SMALLEST TIME PERIOD ON TOP AND LARGEST TIME PERIOD AT THE + BOTTOM. THIS ENSURES PROPER FUNCTIONALITY. + WARNING: TIMING STARTS WHEN THE FIRST PERMITTED REQUEST GOES THROUGH. THIS MEANS THAT, IF YOU HAVE A PER-DAY LIMIT, + AND YOU START MAKING REQUESTS AT 11:00 PM AND EXHAUST YOUR LIMIT AT 11:59 PM, YOUR LIMIT WILL BE REPLENISHED AT + 11:00 PM OF THE NEXT DAY, NOT AT 12:00 AM. + :param attr_name: The name of the variable that holds the instance of 'AsyncRedisCache'. Should be available in the + context of 'current_app'. + :param rate_limit: The number of requests per unit time. + :param seconds: The time period in which the rate limit is to be applied. + :param message: The custom message to respond with. + :param header_keys: The keys in the header to consider when apply rate limits (like 'sessionToken'). + :param data_keys: The keys in the header to consider when apply rate limits (like 'sessionToken'). + :param allow_if_exception: In case Redis is unresponsive, would you prefer allowing the request to pass through or + would you prefer the request getting blocked. + :param count_for_http_codes: If this is provided, the counter will be incremented only if the response code was one + of these values. If not provided, all requests will be counted. This can be used in cases when you want to count + only when the request was successfully served. + :return: The decorator factory. + """ + + # Param-cleaning: + rate_limit = max(1, rate_limit) + if header_keys is None: header_keys = [] + if data_keys is None: data_keys = [] + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Note down the combination of values requested and the limits prescribed: + params = {"route": str(request.url_rule), "header": {}, "data": {}, "limit": rate_limit, "seconds": seconds} + for key in header_keys: params["header"][key] = kwargs["inbound_headers"].get(key) + for key in data_keys: params["data"][key] = kwargs["inbound_data"].get(key) + + # Now make a unique key from this combination: + params_json = json.to_string(params, no_space = True) + sha256_hash = hashlib.sha256() + sha256_hash.update(params_json.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # We first get the value of the counter: + app_attr = getattr(current_app, attr_name) + counter_value = await app_attr.count(base64_key, value = 1, expiry = seconds) + + # If any exception occurred in getting the count, + # and exceptions haven't been allowed: + if counter_value is None and not allow_if_exception: + response = ResponseModel( + status_code = StatusCodes.RATE_LIMIT_EXCEEDED, + message = "Please contact admin (E)" + ) + + # If the rate-limit has already been crossed, + # or when the counter was not fetched but exceptions are allowed: + elif (counter_value or 0) > rate_limit: + response = ResponseModel( + status_code = StatusCodes.RATE_LIMIT_EXCEEDED, + message = message or ", ".join([ + f"rate limit: {rate_limit} in {humanize.naturaldelta(datetime.timedelta(seconds = seconds))}", + f"this is your {humanize.ordinal(counter_value)} request in the given period" + ]) + ) + + # If the rate-limit hasn't been crossed: + else: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # If we have been told to count only for specific status codes, + # and if the HTTP code of this response is not in the list of codes, we reduce the counter by one: + if count_for_http_codes: + http_code = 200 # ... Default assumption. + if isinstance(response, (list, tuple, set)): http_code = response[1] + elif isinstance(response, ResponseModel): http_code = response.http_code.value + if http_code not in count_for_http_codes: + await app_attr.count(base64_key, value = -1, expiry = seconds) + + # Return the response from the function call: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def handle_cancelled_request(cleanup_func = None, cleanup_coro = None): + + """ + Use this decorator to handle prematurely terminated requests. If your clean-up function needs access to variables, + consider using 'g' to hold data in the scope of the request. + :param cleanup_func: The function to call when the cancelled request needs graceful handling. + :param cleanup_coro: The coroutine to call when the cancelled request needs graceful handling. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + try: + + # we are ready to run the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + # In case the client closes the connection pre-maturely: + except asyncio.CancelledError as exception: + kwargs["decorator_count"] -= 1 + if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) + if cleanup_func is not None: cleanup_func() + if cleanup_coro is not None: await cleanup_coro() + return ResponseModel( + api_version = kwargs.get("api_version"), + status_code = StatusCodes.CLIENT_CLOSED_REQUEST + ).for_quart() + + # We propagate any other kind of exception: + except Exception as exception: + kwargs["decorator_count"] -= 1 + raise + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def handle_failed_request(cleanup_func = None, cleanup_coro = None): + + """ + Use this decorator to handle requests that fail due to exceptions. If your clean-up function needs access to + variables, consider using 'g' to hold data in the scope of the request. THIS DECORATOR IS NOT MEANT TO SUPRESS + EXCEPTIONS. IT IS, INSTEAD, MEANT TO PERFORM CLEAN-UP AND PROPAGATE THE EXCEPTION. + :param cleanup_func: The function to call when the cancelled request needs graceful handling. + :param cleanup_coro: The coroutine to call when the cancelled request needs graceful handling. + :return: The decorator factory. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + try: + + # we are ready to run the function that we are wrapping: + kwargs["decorator_count"] = kwargs.get("decorator_count", 0) + 1 + response = await func(*args, **kwargs) + kwargs["decorator_count"] -= 1 + + # Done here: + if ( + kwargs["decorator_count"] == 0 and + isinstance(response, ResponseModel) + ): response = response.for_quart() + return response + + # In case the client closes the connection pre-maturely: + except Exception as exception: + kwargs["decorator_count"] -= 1 + if hasattr(current_app, "printer"): getattr(current_app, "printer")(exception) + if cleanup_func is not None: cleanup_func() + if cleanup_coro is not None: await cleanup_coro() + raise exception + + return wrapper + + return decorator + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/api/codes.py b/utils_v2/api/codes.py new file mode 100644 index 0000000..5773cc9 --- /dev/null +++ b/utils_v2/api/codes.py @@ -0,0 +1,164 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 12th Sept., 2024 + + OBJECTIVE: + + To maintain all status codes in one place. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +from enum import Enum, unique + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +@unique +class HttpCodes(Enum): + + """ + Commonly used standard HTTP status codes. + Can be sent after an API Call is processed. + Refer to: https://http.dev/status + NOTE: THIS LIST IS NOT EXHAUSTIVE! + """ + + # 1XX - Informational: + CONTINUE = 100 # .............. https://http.dev/102 + SWITCHING_PROTOCOLS = 101 # ... https://http.dev/101 + PROCESSING = 102 # ............ https://http.dev/102 + EARLY_HINTS = 103 # ........... https://http.dev/103 + + # 2XX - Success: + SUCCESS = 200 # .................. https://http.dev/200 + CREATED = 201 # .................. https://http.dev/201 + ACCEPTED = 202 # ................. https://http.dev/202 + NON_AUTHORITATIVE_INFO = 203 # ... https://http.dev/203 + NO_CONTENT = 204 # ............... https://http.dev/204 + RESET_CONTENT = 205 # ............ https://http.dev/205 + PARTIAL_CONTENT = 206 # .......... https://http.dev/206 + MULTI_STATUS = 207 # ............. https://http.dev/207 + ALREADY_REPORTED = 208 # ......... https://http.dev/208 + THIS_IS_FINE = 218 # ............. https://http.dev/218 + IM_USED = 226 # .................. https://http.dev/226 + + # 3XX - Redirection: + MULTIPLE_CHOICES = 300 # ..... https://http.dev/300 + MOVED_PERMANENTLY = 301 # .... https://http.dev/301 + MOVED_TEMPORARILY = 302 # .... https://http.dev/302 + SEE_OTHER = 303 # ............ https://http.dev/303 + NOT_MODIFIED = 304 # ......... https://http.dev/304 + USE_PROXY = 305 # ............ https://http.dev/305 + SWITCH_PROXY = 306 # ......... https://http.dev/306 + TEMPORARY_REDIRECT = 307 # ... https://http.dev/307 + PERMANENT_REDIRECT = 308 # ... https://http.dev/308 + + # 4XX - Client Errors: + BAD_REQUEST = 400 # ..................... https://http.dev/401 + UNAUTHORIZED = 401 # .................... https://http.dev/401 + PAYMENT_REQUIRED = 402 # ................ https://http.dev/402 + FORBIDDEN = 403 # ....................... https://http.dev/403 + NOT_FOUND = 404 # ....................... https://http.dev/404 + METHOD_NOT_ALLOWED = 405 # .............. https://http.dev/405 + NOT_ACCEPTABLE = 406 # .................. https://http.dev/406 + PROXY_AUTH_REQUIRED = 407 # ............. https://http.dev/407 + REQUEST_TIMEOUT = 408 # ................. https://http.dev/408 + CONFLICT = 409 # ........................ https://http.dev/409 + GONE = 410 # ............................ https://http.dev/410 + LENGTH_REQUIRED = 411 # ................. https://http.dev/411 + PRECONDITION_FAILED = 412 # ............. https://http.dev/412 + PAYLOAD_TOO_LARGE = 413 # ............... https://http.dev/413 + URI_TOO_LONG = 414 # .................... https://http.dev/414 + UNSUPPORTED_MEDIA_TYPE = 415 # .......... https://http.dev/415 + PAGE_EXPIRED = 419 # .................... https://http.dev/419 + TOO_MANY_REQUESTS = 429 # ............... https://http.dev/429 + UNAVAILABLE_FOR_LEGAL_REASONS = 451 # ... https://http.dev/451 + INVALID_TOKEN = 498 # ................... https://http.dev/498 + CLIENT_CLOSED_REQUEST = 499 # ........... https://http.dev/499 + + # 5XX - Server Errors: + INTERNAL_SERVER_ERROR = 500 # ........... https://http.dev/500 + NOT_IMPLEMENTED = 501 # ................. https://http.dev/501 + BAD_GATEWAY = 502 # ..................... https://http.dev/502 + SERVICE_UNAVAILABLE = 503 # ............. https://http.dev/503 + GATEWAY_TIMEOUT = 504 # ................. https://http.dev/504 + HTTP_VERSION_NOT_SUPPORTED = 505 # ...... https://http.dev/505 + VARIANT_ALSO_NEGOTIATES = 506 # ......... https://http.dev/506 + INSUFFICIENT_STORAGE = 507 # ............ https://http.dev/507 + LOOP_DETECTED = 508 # ................... https://http.dev/508 + BANDWIDTH_LIMIT_EXCEEDED = 509 # ........ https://http.dev/509 + WEB_SERVER_DOWN = 521 # ................. https://http.dev/521 + ORIGIN_IS_UNREACHABLE = 523 # ........... https://http.dev/523 + SERVICE_IS_OVERLOADED = 529 # ........... https://http.dev/529 + NETWORK_READ_TIMEOUT_ERROR = 598 # ...... https://http.dev/598 + NETWORK_CONNECT_TIMEOUT_ERROR = 599 # ... https://http.dev/599 + + +# --------------------------------------------------------------------------------------------------------------------- + + +@unique +class StatusCodes(Enum): + + """ + To be used internally within the context of your service. Customize these to match your service. + The format is: (SUCCESS_INDICATOR, INTERNAL_NUMERIC_CODE, HTTP_CODE) + Example: (True, 1, 200) + """ + + # Legacy Codes: + OK = (True, 1, HttpCodes.SUCCESS.value) + FAILED = (False, 0, HttpCodes.INTERNAL_SERVER_ERROR.value) + PARTIAL_SUCCESS = (True, 2, HttpCodes.PARTIAL_CONTENT.value) + PARTIAL_FAILURE = (False, 3, HttpCodes.PARTIAL_CONTENT.value) + + # Authentication Codes: + LOGGED_IN_SUCCESSFULLY = (True, 200, HttpCodes.SUCCESS.value) + LOGIN_FAILED = (False, 201, HttpCodes.UNAUTHORIZED.value) + INVALID_SESSION_TOKEN = (False, 202, HttpCodes.UNAUTHORIZED.value) + AUTHENTICATION_DETAILS_INCOMPLETE = (False, 203, HttpCodes.BAD_REQUEST.value) + + # Authorization Codes: + AUTHORIZED_SUCCESSFULLY = (True, 300, HttpCodes.SUCCESS.value) + NOT_ALLOWED = (False, 301, HttpCodes.FORBIDDEN.value) + AUTHORIZATION_DETAILS_INCOMPLETE = (False, 302, HttpCodes.BAD_REQUEST.value) + AUTHORIZATION_FAILED = (False, 303, HttpCodes.UNAUTHORIZED.value) + + # General failures: + DOWN_FOR_MAINTENANCE = (False, 800, HttpCodes.SERVICE_UNAVAILABLE.value) + UNKNOWN_ERROR = (False, 801, HttpCodes.INTERNAL_SERVER_ERROR.value) + DATA_INCOMPLETE = (False, 802, HttpCodes.BAD_REQUEST.value) + HEADERS_INCOMPLETE = (False, 803, HttpCodes.BAD_REQUEST.value) + FILES_MISSING = (False, 804, HttpCodes.BAD_REQUEST.value) + CLIENT_CLOSED_REQUEST = (False, 805, HttpCodes.CLIENT_CLOSED_REQUEST.value) + + # Validation failure: + DATA_VALIDATION_FAILURE = (False, 900, HttpCodes.BAD_REQUEST.value) + RATE_LIMIT_EXCEEDED = (False, 901, HttpCodes.TOO_MANY_REQUESTS.value) diff --git a/utils_v2/api/log.py b/utils_v2/api/log.py new file mode 100644 index 0000000..49ccd77 --- /dev/null +++ b/utils_v2/api/log.py @@ -0,0 +1,139 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 12th Sept., 2024 + + OBJECTIVE: + + To have a structure to the response sent from the API calls. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# System-level activities: +import distro +import socket +import platform + +# For data-modelling: +from pydantic import BaseModel, Field +from typing import Any, Optional, List, Literal + +# To work with date and time: +import datetime + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# Info for logging that will stay constant during runtime: +SERVER_HOSTNAME = str(socket.gethostname()) +PLATFORM_INFO = platform.uname() +HOST_OS = str(distro.name(True)) +HOST_CPU = f"{PLATFORM_INFO.processor} ({PLATFORM_INFO.machine})" + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class APILogModel(BaseModel): + + # To identify the machine the code is running on. + # DO NOT MODIFY THESE: + hostname: str = SERVER_HOSTNAME + os: str = HOST_OS + cpu: str = HOST_CPU + # Can modify these: + pid: Optional[Any] = None + ppid: Optional[Any] = None + + # To identify the project and actions: + project: Optional[str] = None + log: str + operation: str + apiVer: Optional[str] = None + logId: Optional[str] = None + logChain: Optional[str] = None + + # Timing metrics: + ts: datetime.datetime + tat: float + cpuTime: float + + # To understand the request that came in: + sessionInfo: Optional[Any] = None + method: Optional[str] = None + url: Optional[str] = None + route: Optional[str] = None + headers: Optional[Any] = None + data: Optional[Any] = None + files: Optional[Any] = None + + # To understand the output that went out: + exception: Optional[Any] = None + response: Optional[Any] = None + httpCode: Optional[int] = None + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + my_log = APILogModel( + log = "internal" + ) + + print(my_log) diff --git a/utils_v2/api/response.py b/utils_v2/api/response.py new file mode 100644 index 0000000..02e3d72 --- /dev/null +++ b/utils_v2/api/response.py @@ -0,0 +1,133 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 12th Sept., 2024 + + OBJECTIVE: + + To have a structure to the response sent from the API calls. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# For data-modelling: +from pydantic import BaseModel +from typing import Any, Optional, List + +# My utils: +from utils_v2.api.codes import StatusCodes, HttpCodes + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class ResponseModel(BaseModel): + + """ + A model for how the response should be when developing API endpoints. + """ + + # The fields that you want in your response: + status_code: StatusCodes + message: Optional[str | List] = None + data: Optional[Any] = None + seconds: Optional[float | int] = None + log_id: Optional[str] = None + http_code: Optional[HttpCodes] = None + api_version: Optional[str] = None + + def for_quart(self): + + """ + Call this when you are using either Flask or Quart as your framework. + :return: The output as expected by Flask and Quart. + """ + + # Construct the basic structure: + response_dict = { + "status": 1 if self.status_code.value[0] else 0, + "code": self.status_code.value[1], + "message": self.message or self.status_code.name.replace("_", " ").lower(), + "data": self.data, + "apiVer": self.api_version + } + + # Now add the additional fields: + if self.seconds is not None: response_dict["seconds"] = self.seconds + if self.log_id is not None: response_dict["logId"] = self.log_id + + # Figure out the HTTP code: + response_http_code = self.http_code.value if self.http_code is not None else self.status_code.value[2] + + # Done here: + return response_dict, response_http_code + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + my_response = ResponseModel( + status_code = StatusCodes.RATE_LIMIT_EXCEEDED + ) + my_response.log_id = "abc123" + + print(my_response.for_quart()) diff --git a/utils_v2/cache/__init__.py b/utils_v2/cache/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/cache/__pycache__/__init__.cpython-310.pyc b/utils_v2/cache/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..b2886ea Binary files /dev/null and b/utils_v2/cache/__pycache__/__init__.cpython-310.pyc differ diff --git a/utils_v2/cache/__pycache__/async_redis_cache.cpython-310.pyc b/utils_v2/cache/__pycache__/async_redis_cache.cpython-310.pyc new file mode 100644 index 0000000..6a0d7d1 Binary files /dev/null and b/utils_v2/cache/__pycache__/async_redis_cache.cpython-310.pyc differ diff --git a/utils_v2/cache/async_redis_cache.py b/utils_v2/cache/async_redis_cache.py new file mode 100644 index 0000000..a0ab8c0 --- /dev/null +++ b/utils_v2/cache/async_redis_cache.py @@ -0,0 +1,497 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 11th April, 2024 + + OBJECTIVE: + + To provide an easy way to work with '.json' files in an async manner. + + REFERENCES: + + 1) https://www.w3schools.com/python/python_json.asp + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To use redis: +import redis.asyncio as redis + +# Other utils: +from utils_v2.string import json + +# For debugging: +from icecream import IceCreamDebugger + +# To make a decorator: +from functools import wraps + +# For hashing and shortening the hash: +import hashlib +import base64 + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** WRAPPERS *** +# ***** **** +# ***************************************************************************************************************** + + +def cache_it(cache = None, expiry = 120): + + """ + This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there. + If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the + cache instead of going through the whole function again. + :param cache: The instance of 'AsyncRedisCache'. + :param expiry: The time in seconds after which the cached data must be cleared. + :return: The decorator that automatically caches your data. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # We first use the name of the function and the inputs given to it to generate a key for Redis with the + # simple hashing and shortening by way of base64 strings: + inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs) + sha256_hash = hashlib.sha256() + sha256_hash.update(inputs_given.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Now we check if we have the value in cache: + try: response = await cache.get(base64_key, raise_exception = True) + + # If the key doesn't exist, we pass through the function and store the results. + except: + response = await func(*args, **kwargs) + await cache.set(key = base64_key, value = response, expiry = expiry) + + # Return the response from the wrapped function. + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def cache_class_methods(attr_name, expiry = 120): + + """ + This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there. + If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the + cache instead of going through the whole function again. + :param attr_name: The name of the variable that has an instance of "AsyncRedisCache". + :param expiry: The time in seconds after which the cached data must be cleared. + :return: The decorator that automatically caches your data. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(self, *args, **kwargs): + + # Get the cache object first: + cache_obj = getattr(self, attr_name) + + # We first use the name of the function and the inputs given to it to generate a key for Redis with the + # simple hashing and shortening by way of base64 strings: + inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs) + sha256_hash = hashlib.sha256() + sha256_hash.update(inputs_given.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Now we check if we have the value in cache: + try: response = await cache_obj.get(base64_key, raise_exception = True) + + # If the key doesn't exist, we pass through the function and store the results. + except: + response = await func(self, *args, **kwargs) + await cache_obj.set(key = base64_key, value = response, expiry = expiry) + + # Return the response from the wrapped function. + return response + + return wrapper + + return decorator + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncRedisCache: + + # Handle datatypes when setting values to cache: + __set_converters = { + "set": lambda x: list(x), + "tuple": lambda x: list(x), + "complex": lambda x: {"r": x.real, "i": x.imag} + } + + # Handle datatypes when getting from cache: + __get_converters = { + "set": lambda x: set(x), + "tuple": lambda x: tuple(x), + "complex": lambda x: complex(x["r"], x["i"]) + } + + def __init__( + self, + connection_string, + ping_counter = 1_000, + debug = False, + debug_prefix = "R-Cache | " + ): + + """ + Implements a simple cache in Redis which holds and returns all native datatypes like ints, floats, bools, + strings, dicts, lists, sets, and tuples :) + :param connection_string: The connection URL for connecting to Redis. + :param ping_counter: The number of requests to Redis after which you want to ping to ensure connection. + :param debug: Whether, or not, you want to show debugging messages from the start. + :param debug_prefix: The prefix text to show with the debugging messages. + """ + + # Note down the configuration: + self.__client = None + self.__ping_counter = ping_counter + self.__requests_since_last_ping = 0 + self.__connection_string = connection_string + + # For debugging: + self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) + if not debug: self.__printer.disable() + + def enable_debug(self): + self.__printer.enable() + + def disable_debug(self): + self.__printer.disable() + + @staticmethod + def make_key(*args, **kwargs): + + """ + Generates a key by hashing the args and kwargs sent to it. Can be useful to generate a predictable key. If the + inputs stay the same, the output stays the same. + :param args: Any number of args that you would like to use to generate the key. + :param kwargs: Any number of kwargs that you would like to use to generate the key. + :return: A string that can be used as a key to store values on Redis. + """ + + # Take everything into one plain text string: + plain_text = "".join(str(a) for a in args) + plain_text += json.to_string(kwargs, no_space = True) + + # Hash the plain text value, and create a key from it: + sha256_hash = hashlib.sha256() + sha256_hash.update(plain_text.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Done here: + return base64_key + + def add_converter( + self, + type_name, + set_converter_func, + get_converter_func, + ): + + """ + To add custom data converters to handle new values. + RULE: Each of the converter functions must take in exactly on argument. + :param type_name: The name of the type of the object to handle. HINT: type(obj).__name__. + :param set_converter_func: A converter function that converts from the custom datatype to a datatype that Redis + can work with. Try converting to an bool, int, float, str, list or dict. + :param get_converter_func: The inverse of the set converter. This function will be used to convert from a + datatype that Redis can work with to the custom datatype. + :return: None. + """ + + self.__set_converters[type_name] = lambda x: set_converter_func(x) + self.__get_converters[type_name] = lambda x: get_converter_func(x) + + async def connect(self): + + """ + To make an asynchronous connection request to the Redis server to establish a connection. + :return: True if connected, else False. + """ + + try: + self.__client = redis.from_url( + self.__connection_string + ) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def disconnect(self): + + """ + Close the connection to the Redis server. + :return: True if successful, else False. + """ + + if self.__client is not None: + try: await self.__client.close() + except Exception as exception: + self.__printer(exception) + return False + return True + return True + + async def ensure_connection(self): + + """ + To ensure that we are connected. We keep pinging the Redis server every once in a while even when connected. + :return: True if connected, else False. + """ + + # If we are not connected, we try to establish a connection: + if self.__client is None: return await self.connect() + + # Else we check if we are connected. If not, we try to connect. + # But we check only once in a while. In the meantime, we assume that we are connected. + elif self.__requests_since_last_ping > self.__ping_counter: + try: + await self.__client.ping() + self.__requests_since_last_ping = 0 + return True + except Exception as exception: + self.__printer(exception) + return await self.connect() + else: + self.__requests_since_last_ping += 1 + return True + + async def set(self, key, value, expiry: float = 120, raise_exception = False): + + """ + Saves some value to the cache. + :param key: The key with which the data will be stored and retrieved. + :param value: The value to store. + :param expiry: The time in seconds after which the data will expire. Must be a positive number. + :param raise_exception: If you want to raise an exception if the process fails. + :return: True if cached, else False. + """ + + await self.ensure_connection() + + try: + + # We handle conversion for unsupported Pythonic datatypes: + original_type = type(value).__name__ + converter = self.__set_converters.get(original_type) + if converter is not None: value = converter(value) + + # Here we actually try to store the data: + response = await self.__client.setex( + key, + int(expiry), + json.to_string({ + "value": value, + "type": original_type + }, indent = 0, no_space = False) + ) + return response + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return False + + async def get(self, key, raise_exception = False, on_fail = None): + + """ + Retrieve the cached value. + :param key: The key with which the data was saved. + :param raise_exception: If you want to raise an exception if the process fails. + :param on_fail: What to return if the process fails due to an exception. + :return: The retrieved data or null if not found. + """ + + await self.ensure_connection() + + try: + data = await self.__client.get(key) + data = data.decode("utf-8") + data = json.from_string(data) + converter = self.__get_converters.get(data["type"]) + if converter is not None: return converter(data["value"]) + return data["value"] + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return on_fail + + async def delete(self, key, raise_exception = False): + + """ + Prematurely delete the value from the cache before it expires. + :param key: The key with which the data was saved. + :param raise_exception: If you want to raise an exception if the process fails. + :return: True if deleted, else False. + """ + + await self.ensure_connection() + + try: + response = await self.__client.delete(key) + return True if response else False + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return False + + async def count(self, key, value: int = 1, expiry: float = None, raise_exception = False): + + """ + To use simple counters. If the counter (identified by the 'key') exists, it will be incremented, else the + counter will be created and the value will be incremented from 0. + :param key: The name of the counter. + :param value: The amount to increment the value by. Send negative values to count backwards. + :param expiry: The time (in seconds) in which the counter expires. Starts from the time the counter is created. + This value has to be an integer. If a float is passed, the value will be rounded off. + :param raise_exception: If you want to raise an exception if the process fails. + :return: The latest value of the counter. Will be null if something went wrong and the exception was suppressed. + """ + + await self.ensure_connection() + + try: + + already_existed = await self.__client.exists(key) + new_value = await self.__client.incrby(key, value) + if expiry and not already_existed: await self.__client.expire(key, int(expiry)) + return new_value + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return None + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import asyncio + + def set_complex(x): + return {"r": x.real, "i": x.imag} + + def get_complex(x): + return complex(x["r"], x["i"]) + + async def main(): + + my_cache = AsyncRedisCache( + connection_string = r"redis://:dc4da94197c843ab6a730113c2b801d9@redis.ditscentre.in/0", + ping_counter = 100, + debug = False + ) + + my_cache.add_converter( + type_name = type(2j).__name__, + set_converter_func = set_complex, + get_converter_func = get_complex + ) + + value = await my_cache.get(key = "6b61afd0-b066-4611-9791-411a30d34624") + print("GET:", value) + print("TYP:", type(value), end = "\n\n") + + # success = await my_cache.set( + # key = "name", + # value = {"first": "John", "last": "Doe"}, + # expiry = 10 + # ) + # print("SET:", success, end = "\n\n") + # + # value = await my_cache.get(key = "name") + # print("GET:", value) + # print("TYP:", type(value), end = "\n\n") + # + # success = await my_cache.delete(key = "name") + # print("DEL:", success, end = "\n\n") + # + # value = await my_cache.get(key = "cnt") + # print("GET:", value) + # print("TYP:", type(value), end = "\n\n") + # + # await my_cache.delete(key = "cnt") + # for _ in range(50): + # await asyncio.sleep(1.0) + # counter = await my_cache.count(key = "cnt", value = 1, expiry = 10) + # print("COUNTER:", counter) + + asyncio.run(main()) diff --git a/utils_v2/cache/async_redis_cache_v2.py b/utils_v2/cache/async_redis_cache_v2.py new file mode 100644 index 0000000..7c82797 --- /dev/null +++ b/utils_v2/cache/async_redis_cache_v2.py @@ -0,0 +1,458 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Saturday, 26th Oct., 2024 + + OBJECTIVE: + + To provide an easy way to cache data for fast access. This version has the change that it can handle custom + serializers by way of dependency injection. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To use redis: +import redis.asyncio as redis + +# Other utils: +from utils_v2.string import json +from utils_v2.serialization.pickle_serializer import PickleSerializer + +# For debugging: +from icecream import IceCreamDebugger + +# To make a decorator: +from functools import wraps + +# For hashing and shortening the hash: +import hashlib +import base64 + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** WRAPPERS *** +# ***** **** +# ***************************************************************************************************************** + + +def cache_it(cache = None, expiry = 120): + + """ + This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there. + If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the + cache instead of going through the whole function again. + :param cache: The instance of 'AsyncRedisCache'. + :param expiry: The time in seconds after which the cached data must be cleared. + :return: The decorator that automatically caches your data. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # We first use the name of the function and the inputs given to it to generate a key for Redis with the + # simple hashing and shortening by way of base64 strings: + inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs) + sha256_hash = hashlib.sha256() + sha256_hash.update(inputs_given.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Now we check if we have the value in cache: + try: response = await cache.get(base64_key, raise_exception = True) + + # If the key doesn't exist, we pass through the function and store the results. + except: + response = await func(*args, **kwargs) + await cache.set(key = base64_key, value = response, expiry = expiry) + + # Return the response from the wrapped function. + return response + + return wrapper + + return decorator + + +# --------------------------------------------------------------------------------------------------------------------- + + +def cache_class_methods(attr_name, expiry = 120): + + """ + This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there. + If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the + cache instead of going through the whole function again. + :param attr_name: The name of the variable that has an instance of "AsyncRedisCache". + :param expiry: The time in seconds after which the cached data must be cleared. + :return: The decorator that automatically caches your data. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(self, *args, **kwargs): + + # Get the cache object first: + cache_obj = getattr(self, attr_name) + + # We first use the name of the function and the inputs given to it to generate a key for Redis with the + # simple hashing and shortening by way of base64 strings: + inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs) + sha256_hash = hashlib.sha256() + sha256_hash.update(inputs_given.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Now we check if we have the value in cache: + try: response = await cache_obj.get(base64_key, raise_exception = True) + + # If the key doesn't exist, we pass through the function and store the results. + except: + response = await func(self, *args, **kwargs) + await cache_obj.set(key = base64_key, value = response, expiry = expiry) + + # Return the response from the wrapped function. + return response + + return wrapper + + return decorator + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncRedisCache: + + def __init__( + self, + connection_string, + serializer = None, + ping_counter = 1_000, + debug = False, + debug_prefix = "R-Cache | " + ): + + """ + Implements a simple cache in Redis which holds and returns all native datatypes like ints, floats, bools, + strings, dicts, lists, sets, and tuples :) + :param connection_string: The connection URL for connecting to Redis. + :param serializer: The serializer to use. + :param ping_counter: The number of requests to Redis after which you want to ping to ensure connection. + :param debug: Whether, or not, you want to show debugging messages from the start. + :param debug_prefix: The prefix text to show with the debugging messages. + """ + + # Note down the configuration: + self.__client = None + self.__serializer = serializer or PickleSerializer() + self.__ping_counter = ping_counter + self.__requests_since_last_ping = 0 + self.__connection_string = connection_string + + # For debugging: + self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) + if not debug: self.__printer.disable() + + def enable_debug(self): + self.__printer.enable() + + def disable_debug(self): + self.__printer.disable() + + @staticmethod + def make_key(*args, **kwargs): + + """ + Generates a key by hashing the args and kwargs sent to it. Can be useful to generate a predictable key. If the + inputs stay the same, the output stays the same. + :param args: Any number of args that you would like to use to generate the key. + :param kwargs: Any number of kwargs that you would like to use to generate the key. + :return: A string that can be used as a key to store values on Redis. + """ + + # Take everything into one plain text string: + plain_text = "".join(str(a) for a in args) + plain_text += json.to_string(kwargs, no_space = True) + + # Hash the plain text value, and create a key from it: + sha256_hash = hashlib.sha256() + sha256_hash.update(plain_text.encode("utf-8")) + hashed_key = sha256_hash.digest() + base64_key = base64.b64encode(hashed_key).decode("utf-8") + + # Done here: + return base64_key + + async def connect(self): + + """ + To make an asynchronous connection request to the Redis server to establish a connection. + :return: True if connected, else False. + """ + + try: + self.__client = redis.from_url( + self.__connection_string + ) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def disconnect(self): + + """ + Close the connection to the Redis server. + :return: True if successful, else False. + """ + + if self.__client is not None: + try: await self.__client.close() + except Exception as exception: + self.__printer(exception) + return False + return True + return True + + async def ensure_connection(self): + + """ + To ensure that we are connected. We keep pinging the Redis server every once in a while even when connected. + :return: True if connected, else False. + """ + + # If we are not connected, we try to establish a connection: + if self.__client is None: return await self.connect() + + # Else we check if we are connected. If not, we try to connect. + # But we check only once in a while. In the meantime, we assume that we are connected. + elif self.__requests_since_last_ping > self.__ping_counter: + try: + await self.__client.ping() + self.__requests_since_last_ping = 0 + return True + except Exception as exception: + self.__printer(exception) + return await self.connect() + else: + self.__requests_since_last_ping += 1 + return True + + async def set(self, key, value, expiry: float = None, raise_exception = False): + + """ + Saves some value to the cache. If an expiry is specified, the data will be deleted after that many seconds. + :param key: The key with which the data will be stored and retrieved. + :param value: The value to store. + :param expiry: The time in seconds after which the data will expire. Must be a positive number. + :param raise_exception: If you want to raise an exception if the process fails. + :return: True if cached, else False. + """ + + await self.ensure_connection() + + try: + + # Here we actually try to set the data: + data = self.__serializer.serialize(value) + if expiry: response = await self.__client.setex(key, int(expiry), data) + else: response = await self.__client.set(key, data) + return response + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return False + + async def get(self, key, raise_exception = False, on_fail = None): + + """ + Retrieve the cached value. + :param key: The key with which the data was saved. + :param raise_exception: If you want to raise an exception if the process fails. + :param on_fail: What to return if the process fails due to an exception. + :return: The retrieved data or null if not found. + """ + + await self.ensure_connection() + + try: + + # Here we try to fetch the data: + data = await self.__client.get(key) + data = self.__serializer.deserialize(data) + return data + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return on_fail + + async def delete(self, key, raise_exception = False): + + """ + Prematurely delete the value from the cache before it expires. + :param key: The key with which the data was saved. + :param raise_exception: If you want to raise an exception if the process fails. + :return: True if deleted, else False. + """ + + await self.ensure_connection() + + try: + + # Try to manually delete the key before expiry: + response = await self.__client.delete(key) + return True if response else False + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return False + + async def count(self, key, value: int = 1, expiry: float = None, raise_exception = False): + + """ + To use simple counters. If the counter (identified by the 'key') exists, it will be incremented, else the + counter will be created and the value will be incremented from 0. + :param key: The name of the counter. + :param value: The amount to increment the value by. Send negative values to count backwards. + :param expiry: The time (in seconds) in which the counter expires. Starts from the time the counter is created. + This value has to be an integer. If a float is passed, the value will be rounded off. + :param raise_exception: If you want to raise an exception if the process fails. + :return: The latest value of the counter. Will be null if something went wrong and the exception was suppressed. + """ + + await self.ensure_connection() + + try: + + already_existed = await self.__client.exists(key) + new_value = await self.__client.incrby(key, value) + if expiry and not already_existed: await self.__client.expire(key, int(expiry)) + return new_value + + except Exception as exception: + self.__printer(exception) + if raise_exception: raise + else: return None + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import asyncio + + def set_complex(x): + return {"r": x.real, "i": x.imag} + + def get_complex(x): + return complex(x["r"], x["i"]) + + async def main(): + + my_cache = AsyncRedisCache( + connection_string = r"redis://:dc4da94197c843ab6a730113c2b801d9@redis.ditscentre.in/0", + ping_counter = 100, + debug = False + ) + + my_cache.add_converter( + type_name = type(2j).__name__, + set_converter_func = set_complex, + get_converter_func = get_complex + ) + + value = await my_cache.get(key = "6b61afd0-b066-4611-9791-411a30d34624") + print("GET:", value) + print("TYP:", type(value), end = "\n\n") + + # success = await my_cache.set( + # key = "name", + # value = {"first": "John", "last": "Doe"}, + # expiry = 10 + # ) + # print("SET:", success, end = "\n\n") + # + # value = await my_cache.get(key = "name") + # print("GET:", value) + # print("TYP:", type(value), end = "\n\n") + # + # success = await my_cache.delete(key = "name") + # print("DEL:", success, end = "\n\n") + # + # value = await my_cache.get(key = "cnt") + # print("GET:", value) + # print("TYP:", type(value), end = "\n\n") + # + # await my_cache.delete(key = "cnt") + # for _ in range(50): + # await asyncio.sleep(1.0) + # counter = await my_cache.count(key = "cnt", value = 1, expiry = 10) + # print("COUNTER:", counter) + + asyncio.run(main()) diff --git a/utils_v2/database/__init__.py b/utils_v2/database/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/database/__pycache__/__init__.cpython-310.pyc b/utils_v2/database/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..346cb74 Binary files /dev/null and b/utils_v2/database/__pycache__/__init__.cpython-310.pyc differ diff --git a/utils_v2/database/__pycache__/async_mongo_v2.cpython-310.pyc b/utils_v2/database/__pycache__/async_mongo_v2.cpython-310.pyc new file mode 100644 index 0000000..9a3b397 Binary files /dev/null and b/utils_v2/database/__pycache__/async_mongo_v2.cpython-310.pyc differ diff --git a/utils_v2/database/__pycache__/async_mysql_v2.cpython-310.pyc b/utils_v2/database/__pycache__/async_mysql_v2.cpython-310.pyc new file mode 100644 index 0000000..a2a82e4 Binary files /dev/null and b/utils_v2/database/__pycache__/async_mysql_v2.cpython-310.pyc differ diff --git a/utils_v2/database/async_firebase.py b/utils_v2/database/async_firebase.py new file mode 100644 index 0000000..a8065a1 --- /dev/null +++ b/utils_v2/database/async_firebase.py @@ -0,0 +1,504 @@ +""" + + AUTHOR: + + Khushal P Soonderji + Bhushan Thakkar + + DATE: + + Friday, 19th April, 2024 + + OBJECTIVE: + + To provide an easy interface to work with Firebase. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +# --- +import sys +sys.path.append(".") +sys.path.append("..") + +# For async operations: +# --- +import asyncio + +# For system-level activities: +# --- +import os + +# My async utils: +# --- +import async_json_utils + +# Firebase: +# --- +import firebase_admin +import firebase_admin.firestore_async +import firebase_admin.auth +import firebase_admin.db + +# For debugging and logging: +# --- +from icecream import IceCreamDebugger + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncFirebase: + + def __init__( + self, + credentials_json_path, + app_name, + max_connections = 5, + debug = True + ): + + # Debugging print: + # --- + self.__printer = IceCreamDebugger(prefix = f"FBase ({app_name[:8]}) | ", includeContext = True) + if not debug: self.__printer.disable() + alert = f"Firebase session starting with max. {max_connections} connections." + self.__printer(alert) + + # Initialize the instance: + # --- + self.__semaphore = asyncio.Semaphore(max_connections) + credentials = firebase_admin.credentials.Certificate(credentials_json_path) + self.__firebase_app = firebase_admin.initialize_app(credentials, name = app_name) + self.__firestore = firebase_admin.firestore_async.client(self.__firebase_app) + + def __del__(self): + firebase_admin.delete_app(self.__firebase_app) + self.__firestore.close() + alert = "Firebase session ended." + self.__printer(alert) + + def __user_to_json(self, firebase_user): + + user_json = { + "uid": firebase_user.uid, + "email": firebase_user.email, + "emailVerified": firebase_user.email_verified, + "displayName": firebase_user.display_name, + "phoneNo": firebase_user.phone_number, + "photoUrl": firebase_user.photo_url, + "customClaims": firebase_user.custom_claims, + "disabled": firebase_user.disabled, + "providerId": firebase_user.provider_id, + "providerData": firebase_user.provider_data, + "tenantId": firebase_user.tenant_id + } + + return user_json + + async def create_custom_token(self, uid): + + async with self.__semaphore: + + try: + custom_token = firebase_admin.auth.create_custom_token( + uid = uid, + app = self.__firebase_app + ).decode("utf-8") + return custom_token + + except Exception as exception: + self.__printer(exception) + return None + + async def authenticate_token(self, token): + + """ + To authenticate the given session token. + :param token: The session token generated by Firebase on a successful sign-in. + :return: Either the retrieved user information or a blank dictionary. + """ + + try: + user_info = firebase_admin.auth.verify_id_token(token, app = self.__firebase_app) + return user_info + except Exception as exception: + self.__printer(exception) + return {} + + async def create_user( + self, + uid, + display_name = None, + password = None, + email = None, + phone_number = None + ): + + """ + Create a new user. + :param uid: The id to identify the user by. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + firebase_admin.auth.create_user( + uid = uid, + app = self.__firebase_app + ) + return True + + except firebase_admin.auth.UidAlreadyExistsError as excp: + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def delete_user(self, uid): + + async with self.__semaphore: + + try: + firebase_admin.auth.delete_user( + uid = uid, + app = self.__firebase_app + ) + return True + + except firebase_admin.auth.UserNotFoundError as excp: + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def get_user(self, uid): + + async with self.__semaphore: + + try: + user = firebase_admin.auth.get_user( + uid = uid, + app = self.__firebase_app + ) + return self.__user_to_json(user) + + except Exception as exception: + self.__printer(exception) + return None + + async def list_users(self, users_per_page, page_token = None): + + """ + To get a list of users. Firebase has a limit of 1000 per call of this API. So we use page-tokens to fetch next + pages of users. + :param users_per_page: How many users you want to list in this API call (Max. 1,000). + :param page_token: To be used in case of pagination. + :return: The list of users and the page token to be used for the next call. Will be None in case of failure. + """ + + async with self.__semaphore: + + try: + users = firebase_admin.auth.list_users( + page_token = page_token, + max_results = min(users_per_page, 1000), + app = self.__firebase_app + ) + next_page_token = users.next_page_token if users.has_next_page else None + users = users.users + users_json = [self.__user_to_json(user) for user in users] + return users_json, next_page_token + + except Exception as exception: + self.__printer(exception) + return None, None + + async def create_document(self, collection_path, document_name, document_data = None): + + """ + To create a new document with the specified data in an existing collection. + :param collection_path: The path of the collection (can be a sub-collection). + :param document_name: The name of the document you want to create. + :param document_data: The data that you want to populate in the document. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.collection( + collection_path + ).add( + document_id = document_name, + document_data = document_data or {} + ) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def delete_document(self, path): + + """ + To delete a document. + :param path: The path of the document. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.document(path).delete() + return True + + except Exception as excp: + print("FIRESTORE DOCUMENT DELETION EXCEPTION:", excp) + return False + + async def set_document(self, path, data): + + """ + To overwrite the data in a document. + :param path: The path of the document. + :param data: The data you want to update as a dictionary. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.document(path).set(data) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def update_document(self, path, data): + + """ + To update the data in a document. Provide only the fields that you want to update. + :param path: The path of the document. + :param data: The data you want to update as a dictionary. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.document(path).update(data) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def get_document(self, path): + + """ + To get the data in a document. + :param path: The path of the document. + :return: The dictionary of data as found in the path specified or None if the operation failed. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.document(path).get() + return snapshot.to_dict() + + except Exception as exception: + self.__printer(exception) + return None + + async def get_document_fields(self, path, fields): + + """ + To get only specific fields (keys) in a document. Like how projections are used in Mongo. + :param path: The path of the document. + :param fields: The list of fields (keys) of the document that you want. + :return: The dictionary of data as found in the path specified. + """ + + async with self.__semaphore: + + try: + if type(fields) is not list: fields = [fields] + snapshot = await self.__firestore.document(path).get(fields) + return snapshot.to_dict() + + except Exception as exception: + self.__printer(exception) + return None + + async def create_collection(self, collection_name, document_name, document_data): + + """ + Creates a new collection in the root of the database. Note that Firebase doesn't allow creating new empty + collections, so we must add one first document in it. + :param collection_name: The name of the collection you want to create. + :param document_name: The name of the first document you want to put in the collection. + :param document_data: The data that you want to put in the first document of the new collection. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.collection( + collection_name + ).add( + document_id = document_name, + document_data = document_data + ) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def create_sub_collection( + self, + document_path, + sub_collection_name, + sub_document_name = None, + sub_document_data = None + ): + + """ + Firebase allows you to create collections inside documents. This method is built for that. Note that Firebase + doesn't allow creating new empty collections, so we must add one first document in it. + :param document_path: The path of the document in which you want to create a new collection. + :param sub_collection_name: The name of the collection you want to create. + :param sub_document_name: The name of the first document you want to put in the collection. + :param sub_document_data: The data that you want to put in the first document of the new collection. + :return: True or False based on the success of the operation. + """ + + async with self.__semaphore: + + try: + snapshot = await self.__firestore.document( + document_path + ).collection( + sub_collection_name + ).add( + document_id = sub_document_name, + document_data = sub_document_data or {} + ) + return True + + except Exception as exception: + self.__printer(exception) + return False + + async def get_collection(self, path): + + """ + Get a whole collection's data. + :param path: The path of the collection. + :return: The collection. + """ + + async with self.__semaphore: + + try: + snapshots = await self.__firestore.collection(path).get() + return {snapshot.id: snapshot.to_dict() for snapshot in snapshots} + + except Exception as exception: + self.__printer(exception) + return None + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import time + from shared import constants + + user_data = { + "firebaseUid": None, + "roles": "ClientAdmin", + "sessionToken": "siddhesh_20240419" + } + + async def test(): + my_firebase = AsyncFirebase( + f"{constants.PROJECT_DIRECTORY}/utils/cred/firebase_certs.json", + "myFire", + max_connections = 5 + ) + print("\n\n---\n\n") + results = await my_firebase.get_document("testCollection/testDoc") + print(async_json_utils.to_json_string(results)) + print("\n\n---\n\n") + results = await my_firebase.get_document_fields("testCollection/testDoc", ["sampleMap"]) + print(async_json_utils.to_json_string(results)) + # await my_firebase.create_sub_collection( + # "myDeepCollection/deepDocId", + # "subCollection3", + # "subDoc", + # {"sub_key": "sub_val"} + # ) + # await my_firebase.create_collection( + # "rootCollection", + # "subDoc", + # {"sub_key": "sub_val"} + # ) + # await my_firebase.delete_document("activeSessions/9d402f1502dfd55a4326fa7fc8e6cb7d") + # print(async_json_utils.to_json_string(await my_firebase.get_collection("myDeepCollection"))) + + start_time = time.time() + asyncio.run(test()) + print(f"FINISHED IN {time.time() - start_time} SECONDS.") diff --git a/utils_v2/database/async_mongo.py b/utils_v2/database/async_mongo.py new file mode 100644 index 0000000..c052fac --- /dev/null +++ b/utils_v2/database/async_mongo.py @@ -0,0 +1,843 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Tuesday, 28th May, 2024 + + OBJECTIVE: + + To have one central place from where all async database connectivity happens. + + REFERENCES: + + 01. https://motor.readthedocs.io/en/stable/ + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# For async behaviour: +import asyncio + +# For datetime handling: +import pytz +import datetime + +# MongoDB: +from motor.motor_asyncio import AsyncIOMotorClient +from bson.objectid import ObjectId +from bson.json_util import dumps, loads + +# For debugging: +from icecream import IceCreamDebugger + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** EXCEPTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +class MongoFindException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo find operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# --------------------------------------------------------------------------------------------------------------------- + + +class MongoInsertException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo insert operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# --------------------------------------------------------------------------------------------------------------------- + + +class MongoUpdateException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo update operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# --------------------------------------------------------------------------------------------------------------------- + + +class MongoReplaceException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo replace operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# --------------------------------------------------------------------------------------------------------------------- + + +class MongoDeleteException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo delete operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# --------------------------------------------------------------------------------------------------------------------- + + +class MongoException(Exception): + + def __init__(self, hint = None, origin = None): + self.__hint = hint + self.__origin = origin + + def __str__(self): + message = "mongo operation failed" + if self.__hint is not None: message = f"{message} ({self.__hint})" + if self.__origin is not None: message = f"{self.__origin} --> {message}" + return message + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncMongo: + + __db = None + __db_name = None + __client = None + + def __init__( + self, + database_name = "myDb", + max_connections = 5, + debug = True, + debug_only_errors = True, + host = "localhost", + port = 27017, + connection_string = None + ): + + # Basic config: + self.__max_connections = max_connections + self.__db_name = database_name + self.__host = host, + self.__port = port + self.__connection_string = connection_string + + # For debugging: + self.__debug_only_errors = debug_only_errors + self.__printer = IceCreamDebugger(prefix = f"Mongo ({self.__db_name}) | ", includeContext = True) + if not debug: self.__printer.disable() + + # Rate/access control: + self.__exclusive_semaphore = asyncio.Semaphore(1) + + def enable_debug(self): + self.__printer.enable() + + def disable_debug(self): + self.__printer.disable() + + @staticmethod + def generate_id(): + + """ + Just generates an '_id' in MongoDB style. + :return: The '_id' in MongoDB style. + """ + + return str(ObjectId()) + + async def connect(self): + try: + if self.__connection_string is None: + self.__client = AsyncIOMotorClient( + self.__host, + self.__port, + maxPoolSize = self.__max_connections, + minPoolSize = self.__max_connections + ) + else: + self.__client = AsyncIOMotorClient( + self.__connection_string, + maxPoolSize = self.__max_connections, + minPoolSize = self.__max_connections + ) + self.__db = self.__client.get_database(self.__db_name) + if not self.__debug_only_errors: + server_info = await self.__client.server_info() + self.__printer(server_info) + except Exception as exception: self.__printer(exception) + + async def ensure_connection(self): + if self.__db is None: + async with self.__exclusive_semaphore: + await self.connect() + + @property + async def client(self): + await self.ensure_connection() + return self.__client + + @staticmethod + def dict_to_dot_notation(input_dict, parent_key = "", separator = "."): + items = [] + for k, v in input_dict.items(): + new_key = f"{parent_key}{separator}{k}" if parent_key else k + if isinstance(v, dict) and v: + items.extend(AsyncMongo.dict_to_dot_notation(v, new_key, separator = separator).items()) + else: + items.append((new_key, v)) + return dict(items) + + @staticmethod + def normalize_date_time(document): + + """ + MongoDB doesn't support timezones. A good strategy would be to convert everything to UTC format and store it. + This method does exactly that. Any datetime object is converted to UTC timezone. If the datetime object was + timezone naive, UTC timezone will be applied to it without changing the time value. + :param document: The document that you want to normalize the date-time in. + :return: The document with normalized datetime. + """ + + if isinstance(document, datetime.datetime): + utc_tz = pytz.timezone("UTC") + if document.tzinfo is None: document = utc_tz.localize(document) + else: document = document.astimezone(utc_tz) + + if type(document) is list: + document = [AsyncMongo.normalize_date_time(item) for item in document] + + if type(document) is dict: + document = { + AsyncMongo.normalize_date_time(k): AsyncMongo.normalize_date_time(v) + for k, v in document.items() + } + + return document + + @staticmethod + def __from_json_string(json_data): + return loads(json_data) + + @staticmethod + def __to_json_string(python_data, indent = 4, default = None): + return dumps(python_data, indent = indent, default = default) + + async def list_indexes( + self, + collection, + session = None, + raise_exception = False + ): + + """ + Lists out the indexes of a collection. + :param collection: The collection whose indexes you want to list out, + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The list of indexes or None if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + indexes = None + + # Try to list the indexes: + try: + responses = await self.__db[collection].list_indexes(session = session).to_list(None) + indexes = [{key: value for key, value in response.items()} for response in responses] + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not indexes: raise MongoException(hint = "list indexes") + return indexes + + async def create_index( + self, + collection, + keys, + options = None, + session = None, + raise_exception = False + ): + + """ + Creates an index on a collection. + :param collection: The collection to create the index on. + :param keys: The keys (and sorting) to implement the index on. + :param options: Additional config. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the execution. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the insertion: + try: + options = options or {} + keys = [(k, v) for k, v in keys.items()] + response = await self.__db[collection].create_index(keys, session = session, **options) + if response: success = True + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not success: raise MongoException(hint = "create index") + return success + + async def insert_one( + self, + collection, + document, + session = None, + raise_exception = False + ): + + """ + Insert data into a collection. + :param collection: The collection you want to feed the data into. + :param document: The data to be stored. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The id of the inserted data, or null if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + inserted_id = None + + # Try to make the insertion: + try: + response = await self.__db[collection].insert_one(document.copy(), session = session) + inserted_id = response.inserted_id + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not inserted_id: raise MongoInsertException(hint = f"{collection}") + return inserted_id + + async def insert_many( + self, + collection, + documents, + session = None, + raise_exception = False + ): + + """ + Insert a lot of data into a collection. + :param collection: The collection you want to feed the data into. + :param documents: The data to be stored. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The id of the inserted data, or null if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + inserted_ids = [] + + # Try to make the insertion: + try: + response = await self.__db[collection].insert_many(documents, session = session) + inserted_ids = response.inserted_ids + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not inserted_ids: raise MongoInsertException(hint = f"{collection}") + return inserted_ids + + async def update_one( + self, + collection, + filter_json = None, + update_json = None, + upsert = False, + session = None, + raise_exception = False + ): + + """ + Update one document. + :param collection: The collection you want to update. + :param filter_json: The selection criteria to locate the document to update. + :param update_json: The values you want to update. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the insertion: + try: + response = await self.__db[collection].update_one( + filter_json, + update_json, + upsert = upsert, + session = session + ) + success = False if response.modified_count == 0 else True + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not success: raise MongoUpdateException(hint = f"{collection}") + return success + + async def update_many( + self, + collection, + filter_json = None, + update_json = None, + upsert = False, + session = None, + raise_exception = False + ): + + """ + Update many documents. + :param collection: The collection you want to update. + :param filter_json: The selection criteria to locate the document to update. + :param update_json: The values you want to update. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + update_count = 0 + + # Try to make the insertion: + try: + response = await self.__db[collection].update_many( + filter_json, + update_json, + upsert = upsert, + session = session + ) + update_count = response.modified_count + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not update_count: raise MongoUpdateException(hint = f"{collection}") + return update_count + + async def count( + self, + collection, + filter_json = None + ): + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + if filter_json is None: filter_json = {} + count = await self.__db[collection].count_documents(filter_json) + except Exception as exception: self.__printer(exception) + + # Done here: + return count + + async def bulk_write( + self, + collection, + requests, + session = None, + raise_exception = False + ): + + """ + To perform various individual operations in one go. You will have to import individual actions like "UpdateOne" + and "InsertMany" from PyMongo and pass them as an array of requests (operations) to this method. + :param collection: The collection you want to run the requests on. + :param requests: The array of requests (operations) to be performed. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + response = await self.__db[collection].bulk_write(requests, session = session) + count = response.modified_count + response.inserted_count + response.upserted_count + response.deleted_count + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not count: raise MongoException(hint = "bulk write") + return count + + async def find_many( + self, + collection, + filter_json, + projections = None, + skip = 0, + limit = None, + sort = None, + session = None, + as_json_string = False, + indent = 4, + raise_exception = False + ): + + """ + Finds one or more records that match the given conditions. + :param collection: The name of the collection to perform the search in. + :param filter_json: The filter criteria. + :param projections: What parts of the matching data you want to fetch. + :param skip: The no. of records to skip before picking next ones. Needed for pagination. + :param limit: The max. no. of records you want to fetch. + :param sort: The sorting rules to apply. + :param session: The session if you need to do this in a transaction. + :param as_json_string: Whether you want it as a JSON string or a Python dict/list. + :param indent: The indentation to use if you want it as a JSON string. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # Assume defaults: + if sort is None: sort = {"_id": -1} + if limit is None: limit = 10 + + # try to query the data: + try: + results = await self.__db[collection].find( + filter_json, + projections, + session = session + ).sort(sort).skip(skip).limit(limit).to_list(None) + if as_json_string: results = self.__to_json_string(results, indent = indent, default = str) + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not results: raise MongoFindException(hint = f"{collection}") + return results + + async def find_one( + self, + collection, + filter_json, + projections = None, + session = None, + as_json_string = False, + indent = 4, + raise_exception = False + ): + + """ + Finds one record that matches the given conditions. + :param collection: The name of the collection to perform the search in. + :param filter_json: The filter criteria. + :param projections: What parts of the matching data you want to fetch. + :param session: The session if you need to do this in a transaction. + :param as_json_string: Whether you want it as a JSON string or a Python dict/list. + :param indent: The indentation to use if you want it as a JSON string. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # try to query the data: + try: + results = await self.__db[collection].find_one(filter_json, projections, session = session) + if as_json_string: results = self.__to_json_string(results, indent = indent, default = str) + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not results: raise MongoFindException(hint = f"{collection}") + return results + + async def replace_one( + self, + collection, + filter_json, + replacement_json, + upsert = False, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete many records. + :param filter_json: The filter criteria. + :param replacement_json: The data to put in place of the existing document. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # try to query the data: + try: + result = await self.__db[collection].replace_one( + filter_json, + replacement_json, + upsert = upsert, + session = session + ) + if result.modified_count or result.upserted_id: success = True + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not success: raise MongoReplaceException(hint = f"{collection}") + return success + + async def delete_one( + self, + collection, + filter_json, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete many records. + :param filter_json: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + result = await self.__db[collection].delete_one(filter_json, session = session) + count = result.deleted_count + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not count: raise MongoDeleteException(hint = f"{collection}") + return count + + async def delete_many( + self, + collection, + filter_json, + session = None, + raise_exception = False + ): + + """ + To delete many documents from a collection. + WARNING: sending {} in the filter would mean deleting ALL the documents. + :param collection: The collection from which you want to delete many records. + :param filter_json: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + result = await self.__db[collection].delete_many(filter_json, session = session) + count = result.deleted_count + except Exception as exception: + self.__printer(exception) + + # Check results and return: + if raise_exception and not count: raise MongoDeleteException(hint = f"{collection}") + return count + + async def aggregate( + self, + collection, + pipeline, + limit = None, + raise_exception = False + ): + + """ + Perform an advance query on the data. + :param collection: The collection to perform the query on. + :param pipeline: The pipeline of actions to take. Must be a list. + :param limit: The max. no. of records to retrieve. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # try to perform the aggregation action: + try: results = await self.__db[collection].aggregate(pipeline).to_list(limit) + except Exception as exception: self.__printer(exception) + + # Check results and return: + if raise_exception and not results: raise MongoException(hint = "aggregation") + return results + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/database/async_mongo_storage.py b/utils_v2/database/async_mongo_storage.py new file mode 100644 index 0000000..94df2d8 --- /dev/null +++ b/utils_v2/database/async_mongo_storage.py @@ -0,0 +1,358 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Tuesday, 28th May, 2024 + + OBJECTIVE: + + To have one central place from where all async database connectivity happens. + + REFERENCES: + + 01. https://motor.readthedocs.io/en/stable/ + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +# --- +import sys +sys.path.append(".") +sys.path.append("..") + +# For system-level activity: +import io + +# For async behaviour: +import asyncio + +# MongoDB for File Storage: +from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket +from bson.objectid import ObjectId + +# For debugging: +from icecream import IceCreamDebugger + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncMongoStorage: + + __db = None + __db_name = None + __client = None + __store = None + + def __init__( + self, + connection_string = None, + max_connections = 5, + host_name = "localhost", + port = 27017, + database_name = "fileStore", + debug = True, + debug_prefix = "GridFS (M) | ", + debug_only_errors = True + ): + + # Database Initialization: + self.__host_name = host_name + self.__port = port + self.__db_name = database_name + self.__connection_string = connection_string + self.__max_connections = max_connections + + # Debugging: + self.__debug_only_errors = debug_only_errors + self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) + if not debug: self.__printer.disable() + + @staticmethod + def generate_id(): + + """ + Just generates an '_id' in MongoDB style. + :return: The '_id' in MongoDB style. + """ + + return str(ObjectId()) + + async def connect(self): + + """ + Initialize the database connection. + :return: Nothing. + """ + + if self.__connection_string is None: + self.__client = AsyncIOMotorClient( + self.__host_name, + self.__port, + maxPoolSize = self.__max_connections, + minPoolSize = self.__max_connections + ) + else: + self.__client = AsyncIOMotorClient( + self.__connection_string, + maxPoolSize = self.__max_connections, + minPoolSize = self.__max_connections + ) + self.__db = self.__client.get_database(self.__db_name) + self.__store = AsyncIOMotorGridFSBucket(self.__db) + + @property + def fs(self): + + """ + To access the features that have not been wrapped in this reportlab directly. + This could include things like streaming files chunk-by-chunk. + :return: The file-store instance. + """ + + return self.__store + + async def write_from_memory(self, file_name, file_data, metadata_json = None): + + """ + Save a file (from RAM) to Mongo. Suitable for smaller files (a few MBs max.). + :param file_name: The name of the file. + :param file_data: The data of the file (held in RAM). + :param metadata_json: A JSON of metadata information that can later be used to search files (RECOMMENDED). + :return: The file's id as a string (if it gets saved) or None. + """ + + if self.__store is None: await self.connect() + + file_data.seek(0) + file_size = file_data.__sizeof__() + + file_id = None + try: file_id = await self.__store.upload_from_stream(file_name, file_data, metadata = metadata_json) + except Exception as exception: self.__printer(exception, file_name, file_size, file_id) + if not self.__debug_only_errors: self.__printer(file_name, file_size, file_id) + return str(file_id) + + async def read_to_memory(self, file_identifier, by_id = True): + + """ + To retrieve a file (in RAM) based on the provided identifier. + Suitable for smaller files (a few MBs max.). + :param file_identifier: Either the name or the "_id" of the file. + :param by_id: Set to True if you are fetching by the "_id" of the file. + :return: Either the file (in RAM) or None. + """ + + if self.__store is None: await self.connect() + + file_data = None + + try: + if by_id: grid_out = await self.__store.open_download_stream(ObjectId(file_identifier)) + else: grid_out = await self.__store.open_download_stream_by_name(file_identifier) + file_data = io.BytesIO(await grid_out.read()) + file_data.seek(0) + + except Exception as exception: + file_data = None + self.__printer(exception, file_identifier, by_id) + + if not self.__debug_only_errors: self.__printer(file_identifier, by_id) + return file_data + + async def delete_file_by_id(self, file_id): + + """ + Tries to delete one file by the id. + :param file_id: The id of the file in the database. + :return: True or False based on the success of the operation. + """ + + if self.__store is None: await self.connect() + + deleted = False + + try: + response = await self.__store.delete(file_id = ObjectId(file_id)) + deleted = True + except Exception as exception: + self.__printer(exception, file_id, deleted) + + return deleted + + def __format_metadata_json(self, metadata_json): + + """ + NOTE: ONLY USE WHEN SEARCHING FILES BY METADATA. + MongoDB expects dot-notation while searching for files by the metadata. We are making a function to search + files assuming that the conditions are to be applied to the metadata itself. So this function add the + dot-notation to the right places to conduct a successful search. + :param metadata_json: The JSON to format. + :return: The formatted JSON that has the right dot-notation. + """ + + formatted_metadata_json = {} + + for key, value in metadata_json.items(): + if not key.startswith("$"): key = f"metadata.{key}" + else: + if type(value) is dict: value = self.__format_metadata_json(value) + if type(value) is list: value = [self.__format_metadata_json(item) for item in value] + formatted_metadata_json[key] = value + + return formatted_metadata_json + + async def find_file_by_metadata(self, metadata_json, limit = None, skip = None, sort = None): + + """ + This method only lists the files that match the criteria mentioned in the metadata JSON. + :param metadata_json: The JSON that describes what you want to find. + :param limit: Max. no. of records to retrieve. + :param skip: No. of starting results to skip. Useful for pagination. + :param sort: The sorting conditions to follow. + :return: A list of (JSONs of) files that match the conditions. The list can be empty. + """ + + if self.__store is None: await self.connect() + + files_list = [] + + try: + limit = limit or 10 + skip = skip or 0 + sort = {"_id": -1} if not isinstance(sort, dict) else sort + formatted_metadata_json = self.__format_metadata_json(metadata_json) + return await self.__store.find( + formatted_metadata_json + ).sort(sort).skip(skip).limit(limit).to_list(None) + + except Exception as exception: self.__printer(exception, metadata_json, len(files_list)) + + if not self.__debug_only_errors: self.__printer(metadata_json, len(files_list)) + return files_list + + async def find_file_by_id(self, file_id): + + """ + This method allows you to get the file's info from the id of the file. + :param file_id: The id that was assigned by Mongo during upload. + :return: The file's info or None if the file doesn't exist. + """ + + if self.__store is None: await self.connect() + + file_info = None + + try: + formatted_metadata_json = {"_id": ObjectId(file_id)} + file_info = (await self.__store.find(formatted_metadata_json).to_list(1))[0] + + except Exception as exception: self.__printer(exception, file_id, file_info) + + return file_info + + async def get_file_name(self, file_id): + + """ + Returns the file name if the id of the file is known. + :param file_id: The id of the file as assigned by MongoDB when the file was stored. + :return: The file's name (if it exists), or None. + """ + + # Ensure that we are connected: + if self.__store is None: await self.connect() + + # Ensure that the input given is of 'ObjectId' type: + if type(file_id) is not ObjectId: file_id = ObjectId(str(file_id)) + + # Fetch and return the file name: + files_list = await self.__store.find( + {"_id": file_id}, + {"filename": True} + ).sort({"_id": -1}).limit(1).to_list(None) + try: file_name = files_list[0]["filename"] + except: file_name = None + return file_name + + async def aggregate( + self, + collection, + pipeline, + limit = None + ): + + """ + Perform an advance query on the data. + :param collection: The collection to perform the query on. + :param pipeline: The pipeline of actions to take. Must be a list. + :param limit: The max. no. of records to retrieve. + :return: The array of matching records or null if there was an exception. + """ + + if self.__store is None: await self.connect() + + results = None + try: results = await self.__db[collection].aggregate(pipeline).to_list(limit) + except Exception as exception: self.__printer(exception) + return results + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/database/async_mongo_v2.py b/utils_v2/database/async_mongo_v2.py new file mode 100644 index 0000000..e5311c4 --- /dev/null +++ b/utils_v2/database/async_mongo_v2.py @@ -0,0 +1,1736 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Original: Tuesday, 28th May, 2024 + Modified: Tuesday, 17th Sept., 2024 + + OBJECTIVE: + + To have one central place from where all async MongoDB activity happens. + + REFERENCES: + + 01. https://motor.readthedocs.io/en/stable/ + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# For system-level activity: +import io + +# For async behaviour: +import asyncio + +# MongoDB for File Storage: +from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket +from bson.objectid import ObjectId +from bson.json_util import dumps, loads + +# My utils: +from utils_v2.string import json + +# For datetime handling: +import pytz +import datetime + +# For debugging: +from icecream import IceCreamDebugger + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** EXCEPTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncMongoBase: + + def __init__( + self, + connection_string = None, + max_connections = 5, + host_name = "localhost", + port = 27017, + database_name = "myDb", + debug = True, + debug_prefix = "Mongo | ", + debug_only_errors = True + ): + + # Basic variables that will be needed later: + self._client = None + self._db = None + self._fs = None + + # Database Initialization: + self._host_name = host_name + self._port = port + self._db_name = database_name + self._connection_string = connection_string + self._max_connections = max_connections + + # Debugging: + self._debug_only_errors = debug_only_errors + self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True) + if not debug: self._printer.disable() + + def enable_debug(self): + self._printer.enable() + + def disable_debug(self): + self._printer.disable() + + async def connect(self): + + """ + Initialize the database connection. + :return: None. + """ + + self._printer("Connecting...") + + # If a connection string is given, + # give preference to it: + if self._connection_string is None: + self._client = AsyncIOMotorClient( + self._host_name, + self._port, + maxPoolSize = self._max_connections, + minPoolSize = self._max_connections + ) + + # In the absense of a connection string, + # try to connect through the other credentials: + else: + self._client = AsyncIOMotorClient( + self._connection_string, + maxPoolSize = self._max_connections, + minPoolSize = self._max_connections + ) + + # Debugging print: + if not self._debug_only_errors: + server_info = await self._client.server_info() + self._printer(server_info) + + # Now we connect to the database and the file-store: + self._db = self._client.get_database(self._db_name) + self._fs = AsyncIOMotorGridFSBucket(self._db) + + async def ensure_connection(self): + + """ + Call this at the start of every function to automatically connect to the database when the user of the library + forgets to explicitly connect to the database. + :return: None. + """ + + if self._client is None: await self.connect() + + @staticmethod + def generate_id(as_str = False): + + """ + Just generates an '_id' in MongoDB style. + :param as_str: Set to True to convert the generated id to a string; and False to receive it as an instance of + 'ObjectId'. This is useful when you need to pre-assign ids to files. + :return: The '_id' in MongoDB style. + """ + + generated_id = ObjectId() + return str(generated_id) if as_str else generated_id + + @property + async def client(self): + + """ + returns the client to perform activities that have not been explicitly wrapped in the class. + :return: The client's instance. + """ + + await self.ensure_connection() + return self._client + + @property + async def db(self): + + """ + returns the database connector to perform activities that have not been explicitly wrapped in the class. + :return: The database's connection instance. + """ + + await self.ensure_connection() + return self._fs + + @property + async def fs(self): + + """ + returns the file-store to perform activities that have not been explicitly wrapped in the class. + :return: The file-store instance. + """ + + await self.ensure_connection() + return self._fs + + @staticmethod + def from_json_string(json_data): + + """ + Converts from a JSON string to BSON. + :param json_data: The JSON string to convert to BSON. + :return: The BSON interpretation of the input JSON string. + """ + + return loads(json_data) + + @staticmethod + def to_json_string(data, indent = 4, default = None): + + """ + Converts from a BSON to JSON string. + :param data: The input BSON data. + :param indent: The no. of spaces to put into the string for pretty print. + :param default: The default function to apply to data that cannot be converted directly. + :return: The JSON string from the input data. + """ + + return dumps(data, indent = indent, default = default) + + @staticmethod + def dict_to_dot_notation(input_dict, pk = "", s = "."): + + """ + Converts an input dict to dot notation format. Can be used as a utility to perform searches. + :param input_dict: The dict that you want to convert to dot notation. + :param pk: Parent Key. DO NOT TOUCH (meant to be used during recursion). + :param s: Separator. DO NOT TOUCH (meant to be used during recursion). + :return: The dot notation representation of the input dict. + """ + + items = [] + for k, v in input_dict.items(): + new_key = f"{pk}{s}{k}" if pk else k + if isinstance(v, dict) and v: items.extend(AsyncMongoBase.dict_to_dot_notation(v, new_key, s = s).items()) + else: items.append((new_key, v)) + return dict(items) + + @staticmethod + def normalize_date_time(document): + + """ + MongoDB doesn't support timezones. A good strategy would be to convert everything to UTC format and store it. + This method does exactly that. Any datetime object is converted to UTC timezone. If the datetime object was + timezone naive, UTC timezone will be applied to it without changing the time value. + :param document: The document that you want to normalize the date-time in. + :return: The document with normalized datetime. + """ + + if isinstance(document, datetime.datetime): + utc_tz = pytz.timezone("UTC") + if document.tzinfo is None: document = utc_tz.localize(document) + else: document = document.astimezone(utc_tz) + + if type(document) is list: + document = [AsyncMongoBase.normalize_date_time(item) for item in document] + + if type(document) is dict: + document = { + AsyncMongoBase.normalize_date_time(k): AsyncMongoBase.normalize_date_time(v) + for k, v in document.items() + } + + return document + + @staticmethod + def read_to_ram(file_path): + + """ + Reads a file into a BytesIO object in RAM. + :param file_path: The path to the file on disk. + :return: The file in a BytesIO object. + """ + + with open(file_path, "rb") as file: file_data = file.read() + file_in_ram = io.BytesIO(file_data) + file_in_ram.seek(0) + return file_in_ram + + +# --------------------------------------------------------------------------------------------------------------------- + + +class AsyncMongo(AsyncMongoBase): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + async def list_indexes( + self, + collection, + session = None, + raise_exception = False + ): + + """ + Lists out the indexes of a collection. + :param collection: The collection whose indexes you want to list out, + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The list of indexes or None if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + indexes = None + + # Try to list the indexes: + try: + responses = await self._db[collection].list_indexes(session = session).to_list(None) + indexes = [{key: value for key, value in response.items()} for response in responses] + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return indexes + + async def create_index( + self, + collection, + keys, + options = None, + session = None, + raise_exception = False + ): + + """ + Creates an index on a collection. + :param collection: The collection to create the index on. + :param keys: The keys (and sorting) to implement the index on. + :param options: Additional config. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the execution. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the insertion: + try: + options = options or {} + keys = [(k, v) for k, v in keys.items()] + response = await self._db[collection].create_index(keys, session = session, **options) + if response: success = True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def count( + self, + collection, + filter, + session = None, + raise_exception = False + ): + + """ + Counts the no. of documents that match the given filter condition. + :param collection: The name of the collection to count in. + :param filter: The filter criteria that the documents must satisfy to be counted. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The count if the operation was performed successfully or None is something went wrong and the exception + was suppressed. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = None + + # try to query the data: + try: count = await self._db[collection].count_documents(filter, session = session) + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def insert_one( + self, + collection, + document, + session = None, + raise_exception = False + ): + + """ + Insert data into a collection. + :param collection: The collection you want to feed the data into. + :param document: The data to be stored. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The id of the inserted data, or null if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + inserted_id = None + + # Try to make the insertion: + try: + response = await self._db[collection].insert_one(document.copy(), session = session) + inserted_id = response.inserted_id + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return inserted_id + + async def insert_many( + self, + collection, + documents, + session = None, + raise_exception = False + ): + + """ + Insert a lot of data into a collection. + :param collection: The collection you want to feed the data into. + :param documents: The list of documents to be stored. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The id of the inserted data, or null if the action fails. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + inserted_ids = [] + + # Try to make the insertion: + try: + response = await self._db[collection].insert_many(documents, session = session) + inserted_ids = response.inserted_ids + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return inserted_ids + + async def find_one( + self, + collection, + filter, + projection = None, + session = None, + raise_exception = False + ): + + """ + Finds one record that matches the given conditions. + :param collection: The name of the collection to perform the search in. + :param filter: The filter criteria. + :param projection: What parts of the matching data you want to fetch. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The matching record or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # try to query the data: + try: results = await self._db[collection].find_one(filter, projection, session = session) + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Check results and return: + return results + + async def find_many( + self, + collection, + filter, + projection = None, + skip = 0, + limit = None, + sort = None, + session = None, + raise_exception = False + ): + + """ + Finds one or more records that match the given conditions. + :param collection: The name of the collection to perform the search in. + :param filter: The filter criteria. + :param projection: What parts of the matching data you want to fetch. + :param skip: The no. of records to skip before picking next ones. Needed for pagination. + :param limit: The max. no. of records you want to fetch. + :param sort: The sorting rules to apply. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # Assume defaults: + if sort is None: sort = {"_id": -1} + if limit is None: limit = 10 + + # try to query the data: + try: + results = await self._db[collection].find( + filter, + projection, + sort = sort, + skip = skip, + limit = limit, + session = session + ).to_list(None) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return results + + async def update_one( + self, + collection, + filter = None, + update = None, + upsert = False, + session = None, + raise_exception = False + ): + + """ + Update one document. + :param collection: The collection you want to update. + :param filter: The selection criteria to locate the document to update. + :param update: The values you want to update. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the update: + try: + response = await self._db[collection].update_one( + filter, + update, + upsert = upsert, + session = session + ) + success = True if response.modified_count > 0 or response.upserted_id else False + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def update_many( + self, + collection, + filter = None, + update = None, + upsert = False, + session = None, + raise_exception = False + ): + + """ + Update many documents. + :param collection: The collection you want to update. + :param filter: The selection criteria to locate the document to update. + :param update: The values you want to update. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + update_count = 0 + + # Try to make the insertion: + try: + response = await self._db[collection].update_many( + filter, + update, + upsert = upsert, + session = session + ) + update_count = response.modified_count + response.upserted_count + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return update_count + + async def find_one_and_update( + self, + collection, + filter = None, + update = None, + projection = None, + return_updated = False, + upsert = False, + session = None, + raise_exception = False + ): + + """ + Update one document. + :param collection: The collection you want to update. + :param filter: The selection criteria to locate the document to update. + :param update: The values you want to update. + :param projection: What parts of the matching data you want to fetch. + :param return_updated: To choose whether you want to retrieve the original document or the updated document. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + document = None + + # Try to make the update: + try: + document = await self._db[collection].find_one_and_update( + filter, + update, + projection = projection, + return_document = return_updated, + upsert = upsert, + session = session + ) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return document + + async def replace_one( + self, + collection, + filter, + replacement, + upsert = False, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete many records. + :param filter: The filter criteria. + :param replacement: The data to put in place of the existing document. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # try to query the data: + try: + result = await self._db[collection].replace_one( + filter, + replacement, + upsert = upsert, + session = session + ) + if result.modified_count or result.upserted_id: success = True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def find_one_and_replace( + self, + collection, + filter, + replacement, + projection = None, + return_replaced = False, + upsert = False, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete many records. + :param filter: The filter criteria. + :param replacement: The data to put in place of the existing document. + :param projection: What parts of the matching data you want to fetch. + :param return_replaced: To choose whether you want to retrieve the original document or the updated document. + :param upsert: If you want to insert if the document doesn't already exist. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + document = None + + # try to query the data: + try: + document = await self._db[collection].find_one_and_replace( + filter, + replacement, + projection = projection, + return_document = return_replaced, + upsert = upsert, + session = session + ) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return document + + async def delete_one( + self, + collection, + filter, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete many records. + :param filter: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # try to query the data: + try: + result = await self._db[collection].delete_one(filter, session = session) + success = True if result.deleted_count else False + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def delete_many( + self, + collection, + filter, + session = None, + raise_exception = False + ): + + """ + To delete many documents from a collection. + WARNING: sending {} in the filter would mean deleting ALL the documents. + :param collection: The collection from which you want to delete many records. + :param filter: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + result = await self._db[collection].delete_many(filter, session = session) + count = result.deleted_count + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def find_one_and_delete( + self, + collection, + filter, + projection = None, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param collection: The collection from which you want to delete a record. + :param filter: The filter criteria. + :param projection: What parts of the matching data you want to fetch. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The document that matched your criteria or None. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + document = None + + # try to query the data: + try: + document = await self._db[collection].find_one_and_delete( + filter, + projection = projection, + session = session + ) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return document + + async def bulk_write( + self, + collection, + requests, + session = None, + raise_exception = False + ): + + """ + To perform various individual operations in one go. You will have to import individual actions like "UpdateOne" + and "InsertMany" from PyMongo and pass them as an array of requests (operations) to this method. + :param collection: The collection you want to run the requests on. + :param requests: The array of requests (operations) to be performed. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # try to query the data: + try: + response = await self._db[collection].bulk_write(requests, session = session) + count = response.modified_count + response.inserted_count + response.upserted_count + response.deleted_count + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def aggregate( + self, + collection, + pipeline, + limit = None, + session = None, + raise_exception = False + ): + + """ + Perform an advance query on the data. + :param collection: The collection to perform the query on. + :param pipeline: The pipeline of actions to take. Must be a list. + :param limit: The max. no. of records to retrieve. preferably apply the limit from within the pipeline. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # try to perform the aggregation action: + try: results = await self._db[collection].aggregate(pipeline, session = session).to_list(limit) + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return results + + +# --------------------------------------------------------------------------------------------------------------------- + + +class AsyncMongoStorage(AsyncMongo): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + async def easy_upload( + self, + source, + file_name: str, + file_metadata: dict = None, + file_id = None, + chunk_size: int = None, + session = None, + raise_exception = False + ): + + """ + Easily write one file to MongoDB's GridFS. Ideal for directly uploading small files without having to worry + about any internal mechanisms. + :param source: The actual file, supplied as either a path string or a file-like object, to be written to the + database. + :param file_name: The name of the file as it will be stored on (and retrieved from) GridFS. + :param file_metadata: Any metadata to later search the file by. + :param file_id: Any custom id to be given to the file. TRY STICKING TO THE ID GENERATED BY 'generate_id'. + :param chunk_size: The chunk size (in bytes) to use for storing this file. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The id of the inserted file or None if the upload failed. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + inserted_id = None + + # Try to write the file: + try: + + # Input pre-processing: + if isinstance(source, io.BytesIO): source.seek(0) + elif isinstance(source, str): source = open(source, mode = "rb") + + # If no file id was supplied: + if file_id is None: + inserted_id = await self._fs.upload_from_stream( + filename = file_name, + source = source, + metadata = file_metadata, + chunk_size_bytes = chunk_size, + session = session + ) + + # If a file id was supplied: + else: + await self._fs.upload_from_stream_with_id( + file_id = file_id, + filename = file_name, + source = source, + metadata = file_metadata, + chunk_size_bytes = chunk_size, + session = session + ) + inserted_id = file_id + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return inserted_id + + async def easy_download( + self, + destination, + file_id = None, + file_name = None, + session = None, + raise_exception = False + ): + + """ + Easily read one file from MongoDB's GridFS. Ideal for directly downloading small files without having to worry + about any internal mechanisms. + :param destination: The path on the local disk or a buffer in RAM to save the downloaded data to. + :param file_id: (RECOMMENDED) the id of the save file. + :param file_name: The name of the saved file. NOT RECOMMENDED because you could have many files with the same + name. The best way to tell files apart if from the id. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. The contents of the stored file are written + directly to the destination. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to read the file: + try: + + # Input pre-processing: + if isinstance(destination, str): destination = open(destination, mode = "wb") + + # If a file id is supplied (preferred way): + if file_id is not None: + await self._fs.download_to_stream( + destination = destination, + file_id = file_id, + session = session + ) + if isinstance(destination, io.BytesIO): destination.seek(0) + success = True + + # If a file name is supplied: + elif file_name is not None: + await self._fs.download_to_stream_by_name( + destination = destination, + filename = file_name, + session = session + ) + if isinstance(destination, io.BytesIO): destination.seek(0) + success = True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def count_files( + self, + filter, + session = None, + raise_exception = False + ): + + """ + Counts the no. of documents that match the given filter condition. + :param filter: The filter criteria that the documents must satisfy to be counted. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The count if the operation was performed successfully or None is something went wrong and the exception + was suppressed. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = None + + # try to query the data: + try: count = await self._db["fs.files"].count_documents(filter, session = session) + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def find_one_file( + self, + filter, + projection = None, + session = None, + raise_exception = False + ): + + """ + Finds one record that matches the given conditions. This does NOT return the file itself, it returns the record + that describes the file. + :param filter: The filter criteria. + :param projection: What parts of the matching data you want to fetch. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The matching record or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # try to query the data: + try: results = await self._db["fs.files"].find_one(filter, projection, session = session) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Check results and return: + return results + + async def find_many_files( + self, + filter, + projection = None, + skip = 0, + limit = None, + sort = None, + session = None, + raise_exception = False + ): + + """ + Finds one or more records that match the given conditions. This doesn't return any actual files directly, it + returns the records that describe the files. + :param filter: The filter criteria. + :param projection: What parts of the matching data you want to fetch. + :param skip: The no. of records to skip before picking next ones. Needed for pagination. + :param limit: The max. no. of records you want to fetch. + :param sort: The sorting rules to apply. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The array of matching records or null if there was an exception. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + results = None + + # Assume defaults: + if sort is None: sort = {"_id": -1} + if limit is None: limit = 10 + + # try to query the data: + try: + results = await self._db["fs.files"].find( + filter, + projection, + session = session + ).sort(sort).skip(skip).limit(limit).to_list(None) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return results + + async def update_metadata_for_one_file( + self, + filter: dict, + unset_data: dict = None, + set_data: dict = None, + session = None, + raise_exception = False + ): + + """ + Updates the metadata for one file. + :param filter: The conditions to filter the files by. + :param unset_data: The fields that you want to discard. + :param set_data: The fields that you want to add or update. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the update: + try: + + # Input pre-processing: + update = {} + if unset_data is not None: update["$unset"] = self.dict_to_dot_notation({"metadata": unset_data}) + if set_data is not None: update["$set"] = self.dict_to_dot_notation({"metadata": set_data}) + + # Actual update happens here: + if update: + response = await self._db["fs.files"].update_one( + filter, + update, + session = session + ) + success = False if response.modified_count == 0 else True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def update_metadata_for_many_files( + self, + filter: dict, + unset_data: dict = None, + set_data: dict = None, + session = None, + raise_exception = False + ): + + """ + Updates the metadata for many files. Practically the same as 'update_metadata_for_one' except that the scope of + the modifications is far wider. + :param filter: The conditions to filter the files by. + :param unset_data: The fields that you want to discard. + :param set_data: The fields that you want to add or update. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = None + + # Try to make the update: + try: + + # Input pre-processing: + update = {} + if unset_data is not None: update["$unset"] = self.dict_to_dot_notation({"metadata": unset_data}) + if set_data is not None: update["$set"] = self.dict_to_dot_notation({"metadata": set_data}) + + # Actual update happens here: + if update: + response = await self._db["fs.files"].update_many( + filter, + update, + session = session + ) + count = response.modified_count + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def replace_metadata_for_one_file( + self, + filter: dict, + replacement, + session = None, + raise_exception = False + ): + + """ + Updates the metadata for one file. + :param filter: The conditions to filter the files by. + :param replacement: The new metadata to put inplace of the old one. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to make the update: + try: + + # Actual update happens here: + result = await self._db["fs.files"].update_one( + filter, + {"$set": {"metadata": replacement}}, + upsert = False, + session = session + ) + if result.modified_count or result.upserted_id: success = True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def replace_metadata_for_many_files( + self, + filter: dict, + replacement, + session = None, + raise_exception = False + ): + + """ + Updates the metadata for one file. + :param filter: The conditions to filter the files by. + :param replacement: The new metadata to put inplace of the old one. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = None + + # Try to make the update: + try: + + # Actual update happens here: + response = await self._db["fs.files"].update_many( + filter, + {"$set": {"metadata": replacement}}, + session = session + ) + count = response.modified_count + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return count + + async def delete_file_by_id( + self, + file_id, + session = None, + raise_exception = False + ): + + """ + Deletes one file by the file's id. Deleting files is different from deleting simple documents because, in files, + you also need to clear out the chunks (which are stored in a separate collection). The built-in mechanism of + Motor only provides support to delete by the file's id. + :param file_id: The id of the file that you want to delete. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Try to read the file: + try: + + # Delete the file: + await self._fs.delete(file_id = file_id, session = session) + success = True + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return success + + async def delete_one_file( + self, + filter, + session = None, + raise_exception = False + ): + + """ + To delete one document from a collection. + :param filter: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: True or False based on the success of the operation. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + success = False + + # Find one file that matches the given filter: + file = await self.find_one( + collection = "fs.files", + filter = filter, + session = session, + raise_exception = raise_exception + ) + + # If we have a list of files to work with: + if file: success = await self.delete_file_by_id(file["_id"]) + + # Done here: + return success + + async def delete_many_files( + self, + filter, + session = None, + raise_exception = False + ): + + """ + To delete many documents from a collection. + WARNING: sending {} in the filter would mean deleting ALL the documents. + :param filter: The filter criteria. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The number of files deleted (can be zero) or None if something failed while searching the files. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + count = 0 + + # Find all the files that match the given filter: + files = await self.find_many( + filter = filter, + session = session, + raise_exception = raise_exception + ) + + # If we have a list of files to work with: + if files is not None: + tasks = [self.delete_file_by_id(file["_id"]) for file in files] + results = await asyncio.gather(*tasks) + count = sum(results) + + # Done here: + return count + + async def get_upload_stream( + self, + file_name: str, + file_metadata: dict = None, + file_id = None, + chunk_size: int = None, + session = None, + raise_exception = False + ): + + """ + Returns a GridIn object so that you can perform your own upload using the built-in writing methods. You must use + the 'write' method to write data to the file by passing it either a string of bytes or a file-like object. When + the file has been fully written, you must call the 'close' method to finish the operation. In case you need to + cancel the operation, you can call the 'abort' method to delete all the already written data and stop uploading + new data. + :param file_name: The name of the file as it will be stored on (and retrieved from) GridFS. + :param file_metadata: Any metadata to later search the file by. + :param file_id: Any custom id to be given to the file. TRY STICKING TO THE ID GENERATED BY 'generate_id'. + :param chunk_size: The chunk size (in bytes) to use for storing this file. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The upload stream that implements the 'write', 'close', and 'abort' methods, or None if something + failed. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + stream = None + + # Try to open the upload stream: + try: + + # If no file id was supplied: + if file_id is None: + stream = self._fs.open_upload_stream( + filename = file_name, + metadata = file_metadata, + chunk_size_bytes = chunk_size, + session = session + ) + + # If a file id was supplied: + else: + stream = self._fs.open_upload_stream_with_id( + file_id = file_id, + filename = file_name, + metadata = file_metadata, + chunk_size_bytes = chunk_size, + session = session + ) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return stream + + async def get_download_stream( + self, + file_id = None, + file_name = None, + session = None, + raise_exception = False + ): + + """ + Returns a GridOut object so that you can implement your own download logic using the built-in 'read' method. + Once the reading is done, use the 'close' method to release the resources used by the stream. + :param file_id: (RECOMMENDED) the id of the save file. + :param file_name: The name of the saved file. NOT RECOMMENDED because you could have many files with the same + name. The best way to tell files apart if from the id. + :param session: The session if you need to do this in a transaction. + :param raise_exception: Whether, or not, you want to raise an exception when something fails. + :return: The download stream that implements the 'read' and 'close' methods, or None if something failed. + """ + + # Ensure you are connected: + await self.ensure_connection() + + # Assume failure: + stream = None + + # Try to open a download stream: + try: + + # If a file id is supplied (preferred way): + if file_id is not None: + stream = await self._fs.open_download_stream( + file_id = file_id, + session = session + ) + + # If a file name is supplied: + elif file_name is not None: + stream = await self._fs.open_download_stream_by_name( + filename = file_name, + session = session + ) + + # When something goes wrong: + except Exception as exception: + self._printer(exception) + if raise_exception: raise + + # Done here: + return stream + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + from shared import constants + + async def main(): + + # Create an instance of the database connector: + my_db = AsyncMongo( + connection_string = constants.MONGO_DATA_CONNECTION_STRING, + database_name = constants.MONGO_DATA_DATABASE_NAME, + max_connections = 10, + debug = True + ) + + # Connect to the database: + await my_db.connect() + + # # Get the documents to migrate: + # documents = await my_db.find_many( + # collection = "scriptData", + # filter = {}, + # limit = 50, + # projection = {"_id": False} + # ) + # # print(json.to_string(documents, default = str)) + # + # # Adjust them: + # adjusted_documents = [] + # for document in documents: + # script_id = document.pop("scriptId") + # adjusted_document = { + # "scriptId": script_id, + # "desc": "no desc", + # "content": document + # } + # adjusted_documents.append(adjusted_document) + # print(json.to_string(adjusted_documents, default = str)) + # + # # Insert the adjusted ones to the new collection: + # response = await my_db.insert_many( + # collection = "_scriptData", + # documents = adjusted_documents + # ) + # print("RESPONSE:", response) + + + asyncio.run(main()) diff --git a/utils_v2/database/async_mysql_v2.py b/utils_v2/database/async_mysql_v2.py new file mode 100644 index 0000000..f4cac6e --- /dev/null +++ b/utils_v2/database/async_mysql_v2.py @@ -0,0 +1,424 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Friday, 30th Aug., 2024 + + OBJECTIVE: + + To be able to access SQL-based databases from python in a simple way. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# MySQL Database: +import aiomysql +import decimal + +# For data-crunching: +import pandas as pd + +# For time-keeping: +import time + +# OS-level operations: +import os + +# My utils: +from utils_v2.string import json + +# For async activities: +import asyncio + +# For debugging: +from icecream import IceCreamDebugger +import traceback + +# To work with datatypes: +from typing import List + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncMySQL: + + def __init__( + self, + pool_size, + *args, + **kwargs + ): + + """ + A class to work with SQL-based databases. Originally meant to only invoke stored procedures and retrieve them as + JSON-like structures (list or dict). The format for the results was very specific to our use case for serving + Bicree's requirement. This may not serve your requirement at all. + :param pool_size: The number of connections to maintain n a pool. + :param args: Any arguments to pass. Not used. + :param kwargs: Pass the connection configuration from here. + """ + + # Set up the variables: + self.__args = args + self.__kwargs = kwargs + self.__min_pool_size = 10 + self.__max_pool_size = max(pool_size, self.__min_pool_size) + self.__pool = None + + # Set up the debugging tools: + self.__printer = IceCreamDebugger(prefix = "MySQL | ", includeContext = True) + + def __del__(self): + pass + + async def connect(self): + + """ + Establish a connection and create a pool of connections to call from. + :return: None. + """ + + try: + + self.__kwargs["db"] = self.__kwargs.pop("database") + self.__pool = await aiomysql.create_pool( + minsize = self.__min_pool_size, + maxsize = self.__max_pool_size, + loop = asyncio.get_event_loop(), + **self.__kwargs + ) + + except Exception as exception: + self.__printer(exception) + self.__pool = None + + async def ensure_connection(self): + + """ + Tries to ensure that a connection is present. + Can be called before every function to make sure that our pool is established. + :return: None. + """ + + if self.__pool is None: await self.connect() + + @staticmethod + def __parse_row(row): + + """ + Converts from the custom objects of 'aiomysql' to types that are supported by Python. + :param row: The row from the result. + :return: The parsed row which will have types that are closer to being native to Python.. + """ + + parsed_row = [] + for item in row: + if isinstance(item, decimal.Decimal): parsed_row.append(float(item)) + else: parsed_row.append(item) + return parsed_row + + async def fetch_all(self, cursor): + + # Make a variable to hold all the result sets. + # Needed for when the procedure responds with many "tables": + all_result_sets = [] + + # Iterate over all result sets, + # and process them one-by-one: + while True: + + # Process the current result set: + this_result_set = [] + result = await cursor.fetchall() + if not cursor.description: break + columns = [desc[0] for desc in cursor.description] + for row in result: this_result_set.append(dict(zip(columns, self.__parse_row(row)))) + all_result_sets.append(this_result_set) + + # Move to the next set, + # or break out of the loop if all done: + if not await cursor.nextset(): break + + # Done here: + return all_result_sets + + async def call_procedure(self, procedure_name, procedure_args): + + """ + To call stored procedures and retrieve all the responses. + :param procedure_name: The name of the stored procedure that must be called. + :param procedure_args: The args to be sent to the stored procedure. + :return: The raw result set as received from the database. + """ + + # Make sure we have a connection: + await self.ensure_connection() + + # Make a variable to hold all the result sets. + # Needed for when the procedure responds with many "tables": + all_result_sets = [] + + # Call the procedure and get the results: + async with self.__pool.acquire() as connection: + async with connection.cursor() as cursor: + await cursor.callproc(procedure_name, procedure_args) + all_result_sets = await self.fetch_all(cursor) + + # # Iterate over all result sets, + # # and process them one-by-one: + # while True: + # this_result_set = [] + # result = await cursor.fetchall() + # if not cursor.description: break + # columns = [desc[0] for desc in cursor.description] + # for row in result: this_result_set.append(dict(zip(columns, self.__parse_row(row)))) + # all_result_sets.append(this_result_set) + # await cursor.nextset() + + # Done here: + return all_result_sets + + async def call_procedure_and_get_json( + self, + procedure_name, + procedure_args, + retry_count = 1, + backoff_seconds = 0.5, + backoff_multiplier = 1.1, + return_exception = False + ): + + """ + The method to call when you need to call a stored procedure and retrieve the response as a JSON-like object. + This is custom formatting based on the structure created by Mr. bhushan Thakkar in late April (2024). + :param procedure_name: The name of the stored procedure that must be called. + :param procedure_args: The args to be sent to the stored procedure. + :param retry_count: The max. number of times to try in case one or more attempts fail. + :param backoff_seconds: The time to wait before making the next attempt if the retry count is more than 1. + :param backoff_multiplier: The factor that dictates how much to modify the time delay by when waiting to retry. + :param return_exception: Whether, or not, you would like to return the exception object if something goes wrong. + :return: The formatted response and the exception (if asked for). + """ + + # Note down the start time: + start_ts = time.time() + + # Try to get the data from the database: + results = [] + exception = None + for _ in range(retry_count): + try: results = await self.call_procedure( + procedure_name = procedure_name, + procedure_args = procedure_args, + ) + except Exception as exc: exception = exc + if exception is None: break + await asyncio.sleep(backoff_seconds) + backoff_seconds = backoff_seconds * backoff_multiplier + + # If the results are blank: + if len(results) == 0: + formatted_results = { + "status": 0, + "message": "Please contact admin (NE)" if exception is None else "Please contact admin (E)", + "seconds": time.time() - start_ts, + "data": {} + } + if return_exception: return formatted_results, exception + else: return formatted_results + + # Extract the very basic success or failure indicators: + formatted_results = { + "status": results[0][0]["status"], + "message": results[0][0].get("message", "ok"), + "seconds": 0.0, + "data": {} + } + + # Handle the remaining keys of the zeroth result set: + for key, value in results[0][0].items(): + if key not in formatted_results.keys(): + formatted_results["data"][key] = value + + # Format + for index in range(len(results)): + if index > 0: formatted_results["data"][f"rs{index-1}"] = results[index] + + # Note down the time taken: + formatted_results["seconds"] = time.time() - start_ts + + # Done here: + if return_exception: return formatted_results, exception + else: return formatted_results + + async def execute_one( + self, + query: str, + commit: bool = True, + return_exception: bool = False + ): + + """ + Runs one command / query in SQL. + :param query: The query / command to run. + :param commit: Whether, or not, you would like to commit the execution. + :param return_exception: Whether, or not, you would like to return the exception from this function. + :return: Either just the result or the result and the exception. + """ + + # Make sure we have a connection: + await self.ensure_connection() + + # Start by assuming failure: + rows_affected = None + results = None + excp = None + + try: + + # Get a connection and execute the command: + async with self.__pool.acquire() as connection: + async with connection.cursor() as cursor: + rows_affected = await cursor.execute(query) + results = await self.fetch_all(cursor) + if commit: await connection.commit() + + # SQL-specific errors: + except aiomysql.MySQLError as exception: + self.__printer("SQL Exception", exception) + excp = exception + + # Other errors: + except Exception as exception: + self.__printer("Other Exception", exception) + excp = exception + + # Done here: + if return_exception: return rows_affected, results, excp + else: return rows_affected, results + + async def execute_many( + self, + query: str, + data: List[tuple], + commit: bool = True, + return_exception: bool = False + ): + + """ + Runs many commands / queries in SQL. + Consider the following example: + QUERY: "INSERT INTO pincodeMaster (pincode, city, state) VALUES (%s, %s, %s);" + DATA: [ + ('110001', 'New Delhi', 'Delhi'), + ('500001', 'Hyderabad', 'Telangana'), + ('600001', 'Chennai', 'Tamil Nadu') + ] + :param query: The query / command to run. + :param data: The data to substitute into the query string. + :param commit: Whether, or not, you would like to commit the execution. + :param return_exception: Whether, or not, you would like to return the exception from this function. + :return: Either just the result or the result and the exception. + """ + + # Make sure we have a connection: + await self.ensure_connection() + + # Start by assuming failure: + rows_affected = None + results = None + excp = None + + try: + + # Get a connection and execute the command: + async with self.__pool.acquire() as connection: + async with connection.cursor() as cursor: + rows_affected = await cursor.executemany(query, data) + results = await self.fetch_all(cursor) + if commit: await connection.commit() + + # SQL-specific errors: + except aiomysql.MySQLError as exception: + self.__printer("SQL Exception", exception) + excp = exception + + # Other errors: + except Exception as exception: + self.__printer("Other Exception", exception) + excp = exception + + # Done here: + if return_exception: return rows_affected, results, excp + else: return rows_affected, results + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/dataframe/__init__.py b/utils_v2/dataframe/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/dataframe/splitter.py b/utils_v2/dataframe/splitter.py new file mode 100644 index 0000000..7725905 --- /dev/null +++ b/utils_v2/dataframe/splitter.py @@ -0,0 +1,118 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Tuesday, 22nd Oct., 2024 + + OBJECTIVE: + + To provide an easy way to split dataframes into chunks and divide the workload nto more manageable batches. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + + USAGE EXAMPLE: + + for sub_df in DataFrameSplitter(df, chunk_size = 50): + + # Do your task here: + pass + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class DataFrameSplitter: + + def __init__(self, df, chunk_size): + + """ + Use this to process your dataframe in batches. Useful for when you need to send out alerts at intervals or need + to maintain checkpoints. + :param df: The dataframe to iterate over. + :param chunk_size: The max. no. of rows to pick at once. + """ + + self.df = df + self.row_count = df.shape[0] + self.chunk_size = chunk_size + self.offset = 0 + + def __iter__(self): + self.offset = 0 + return self + + def __next__(self): + if self.offset >= self.row_count: + raise StopIteration + end = self.offset + self.chunk_size + chunk = self.df.iloc[self.offset:end] + self.offset = end + return chunk + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/datetime/__init__.py b/utils_v2/datetime/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/datetime/datetime.py b/utils_v2/datetime/datetime.py new file mode 100644 index 0000000..4659bdd --- /dev/null +++ b/utils_v2/datetime/datetime.py @@ -0,0 +1,278 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Friday, 21st jun, 2024 + + OBJECTIVE: + + To provide an easy way to work with time. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# For date and time handling: +import pytz +from datetime import datetime, time, timedelta +import dateparser + +# To handle date-time objects from a Numpy array and Pandas Dataframe: +import numpy as np +import pandas as pd + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# Date-Time Formats: +DATE_TIME_FORMATS = ( + "%d/%m/%y", + "%d-%b-%y", + "%d-%m-%y", + "%d.%m.%y", + "%d/%m/%Y", + "%d-%b-%Y", + "%d-%m-%Y", + "%d.%m.%Y", + "%d/%b", + "%d%m%Y", + "%Y%m%d", + "%Y-%m-%d %H:%M:%S" +) + +# Useful Timezones: +TIMEZONE_UTC = pytz.timezone("UTC") +TIMEZONE_IST = pytz.timezone("Asia/Kolkata") +TIMEZONE_ET = pytz.timezone("America/New_York") +TIMEZONE_CT = pytz.timezone("America/Chicago") +TIMEZONE_MT = pytz.timezone("America/Denver") +TIMEZONE_PT = pytz.timezone("America/Los_Angeles") +TIMEZONE_JST = pytz.timezone("Asia/Tokyo") +TIMEZONE_CET = pytz.timezone("Europe/Paris") +TIMEZONE_GMT = pytz.timezone("GMT") +TIMEZONE_AEST = pytz.timezone("Australia/Sydney") +TIMEZONE_NZST = pytz.timezone("Pacific/Auckland") +TIMEZONE_CST = pytz.timezone("Asia/Shanghai") +TIMEZONE_KST = pytz.timezone("Asia/Seoul") +TIMEZONE_MSK = pytz.timezone("Europe/Moscow") +TIMEZONE_BRT = pytz.timezone("America/Sao_Paulo") +TIMEZONE_GST = pytz.timezone("Asia/Dubai") +TIMEZONE_SAST = pytz.timezone("Africa/Johannesburg") +TIMEZONE_AST = pytz.timezone("Asia/Riyadh") + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +def translate_date_time_string( + datetime_string, + source_format = None, + destination_format = "%Y-%m-%dT%H:%M:%S" +): + + """ + To convert an input datetime string to a different format. + :param datetime_string: The datetime string to translate. + :param source_format: The current format of the string. If not provided, dateparser will be used. + :param destination_format: The format to convert to. + :return: The converted datetime string. + """ + + try: + if source_format is None: datetime_obj = dateparser.parse(datetime_string) + else: datetime_obj = datetime.strptime(datetime_string, source_format) + return datetime_obj.strftime(destination_format) + except Exception as exception: + return None + + +# --------------------------------------------------------------------------------------------------------------------- + + +def parse_date_time(input_value, timezone = None, date_formats = None): + + """ + To take any kind of input and interpret the datetime from it. + :param input_value: Either a string or an integer or some form of datetime representation. + :param timezone: The timezone to apply to the interpreted datetime. EXISTING TIMEZONE INFO WILL BE OVERWRITTEN. + :param date_formats: The string formats to consider when parsing a string input. + :return: The parsed datetime or null. + """ + + datetime_object = None + date_formats = date_formats or DATE_TIME_FORMATS + + # It could either be in seconds or milliseconds from epoch time's base date (January 1, 1970), + # or it could be days since Microsoft Excel's base date (December 31, 1899). + if isinstance(input_value, (int, float, np.number)) and not np.isnan(input_value): + if input_value > 9999999999.0: datetime_object = datetime.fromtimestamp(input_value / 1000.0) + if input_value > 999999.0: datetime_object = datetime.fromtimestamp(input_value) + else: datetime_object = datetime.fromtimestamp(input_value * 24 * 60 * 60.0) - timedelta(days = 25569) + + # The input can even be a pre-formatted date: + if isinstance(input_value, str): + datetime_object = dateparser.parse( + input_value, + date_formats = date_formats, + settings = { + "DATE_ORDER": "DMY", + "PREFER_DAY_OF_MONTH": "first", + } + ) + + # If the type is a datetime object, then return it as it is: + if isinstance(input_value, datetime): + datetime_object = input_value + + # If the type is the native datetime format of pandas: + if isinstance(input_value, pd._libs.tslibs.timestamps.Timestamp): + datetime_object = input_value.to_pydatetime() + + # Process the timezone: + if datetime_object is not None and timezone is not None: + datetime_object = as_if_timezone(datetime_object, timezone) + + # Done here: + return datetime_object + + +# --------------------------------------------------------------------------------------------------------------------- + + +def get_current_date_time(timezone = None, as_string = False): + + """ + Returns the current time as a datetime object. + :param timezone: The timezone to apply to the returned datetime. + :param as_string: Whether, or not, you want the output as a string. + :return: The datetime object/string representing the current time. + """ + + if timezone is not None and isinstance(timezone, str): timezone = pytz.timezone(timezone) + now = datetime.now(timezone) + return now.isoformat() if as_string else now + + +# --------------------------------------------------------------------------------------------------------------------- + + +def get_current_ist_date_time(as_string = False): + + """ + Gives out the current time in IST timezone. + :param as_string: Whether, or not, you want the output as a string. + :return: The datetime object or string representing the current time. + """ + + return get_current_date_time( + timezone = TIMEZONE_IST, + as_string = as_string + ) + + +# --------------------------------------------------------------------------------------------------------------------- + + +def get_current_utc_date_time(as_string = False): + + """ + Gives out the current time in UTC timezone. + :param as_string: Whether, or not, you want the output as a string. + :return: The datetime object or string representing the current time. + """ + + return get_current_date_time( + timezone = TIMEZONE_UTC, + as_string = as_string + ) + + +# --------------------------------------------------------------------------------------------------------------------- + + +def as_if_timezone(datetime_object, timezone): + + """ + Ignores existing timezone info and applies the intended timezone. + The time stays the same, only the timezone marker changes. + e.g. for IST to UTC: 2024-08-09 00:00:00+05:30 --> 2024-08-09 00:00:00+00:00 + HINT: IT PRETENDS "AS IF" THE TIMEZONE WAS THE INPUT TIMEZONE. + :param datetime_object: The datetime object on which the timezone needs to be applied. + :param timezone: The timezone that needs to be applied. + :return: A timezone-aware datetime object. + """ + + tz_object = pytz.timezone(timezone) if isinstance(timezone, str) else timezone + return tz_object.localize(datetime_object.replace(tzinfo = None)) + + +# --------------------------------------------------------------------------------------------------------------------- + + +def to_timezone(datetime_object, timezone): + + """ + Converts from one timezone to another. The time is adjusted by computing the difference between the two timezones. + NOTE: THIS FUNCTION ASSUMES THE INPUT WAS IN UTC IF THE INPUT WAS TIMEZONE-NAIVE. + e.g. for IST to UTC: 2024-08-09 00:00:00+05:30 --> 2024-08-08 18:30:00+00:00 + :param datetime_object: The datetime object on which the timezone needs to be applied. + :param timezone: The timezone that needs to be applied. + :return: A timezone-aware datetime object. + """ + + if isinstance(timezone, str): timezone = pytz.timezone(timezone) + if datetime_object.tzinfo is None: return datetime_object.replace(tzinfo = TIMEZONE_UTC) + return datetime_object.astimezone(timezone) + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/image/__init__.py b/utils_v2/image/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/image/scanner/__init__.py b/utils_v2/image/scanner/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/image/scanner/scanner.py b/utils_v2/image/scanner/scanner.py new file mode 100644 index 0000000..28fe69b --- /dev/null +++ b/utils_v2/image/scanner/scanner.py @@ -0,0 +1,410 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Thursday, 19th Sept., 2024 + + OBJECTIVE: + + To provide a way to easily scan documents like printouts and visiting cards. + + REFERENCES: + + 01. https://medium.com/@victorolufemi/build-a-document-scanner-with-opencv-ff9f645a4085 + 02. https://github.com/JaidedAI/EasyOCR + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys + +from diffusers.utils.import_utils import candidates + +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import io + +# To work with image processing: +import cv2 +import imutils +import numpy as np + +# To run OCR: +import easyocr + +# To download images from the web: +import requests + +# My utils: +from utils_v2.string import json +from utils_v2.ai.object_detection.yolo import YoloDetect + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class DocumentScanner: + + def __init__( + self, + layout_detection_yolo = None, + whitelisted_yolo_classes = None, + ocr_languages = None + ): + + self.__yolo = YoloDetect( + model_file = layout_detection_yolo, + debug = False + ) + self.__whitelisted_yolo_classes = whitelisted_yolo_classes or [] + self.__ocr_engine = easyocr.Reader(ocr_languages or ["en"]) + pass + + @staticmethod + def open_image(image): + + """ + Opens an image from various sources. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :return: The image as a CV2 object. + """ + + # When the image is provided as a string, + # it could be either as a URL or a path to a local file: + if isinstance(image, str): + if image.startswith("https://") or image.startswith("http://"): + image = io.BytesIO(requests.get(image).content) + else: image = cv2.imread(image) + + # If the image is provided as a io.BytesIO buffer: + if isinstance(image, io.BytesIO): + image.seek(0) + image = np.asarray(bytearray(image.read()), dtype = np.uint8) + image = cv2.imdecode(image, cv2.IMREAD_COLOR) + + # Done here: + return image + + @staticmethod + def polygon_area(points): + + """ + Computes the area occupied by a shape described by the array of points. + Example input: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]] + :return: The area of the shape. + """ + + n = len(points) + area = 0.0 + for i in range(n): + j = (i + 1) % n # next vertex + area += points[i][0] * points[j][1] + area -= points[j][0] * points[i][1] + return abs(area) / 2.0 + + @staticmethod + def show(image, title = "Preview", wait = True): + + """ + Just a quick wrapper to show the image in a window. + :param image: The image that you want to show. + :param title: The title of the window. + :param wait: Set this to True when you want to show the window(s). This is useful when you want to show many + windows at once. Suppose you want to show a lot of windows, you set this to False for all calls except the + very last one. + :return: None. + """ + + cv2.imshow(title, image) + if wait: + cv2.waitKey() + cv2.destroyAllWindows() + + @staticmethod + def enhance(image): + + """ + Applies a form of contrast boost to make the edges more easily visible. + :param image: The OpenCV image that needs to be enhanced. + :return: The enhanced CV2 image. + """ + + return cv2.detailEnhance( + src = image, + sigma_s = 20, + sigma_r = 0.15 + ) + + @staticmethod + def get_edges(image): + + """ + Gets the edges in an image. + :param image: A CV2 image. + :return: The image with the edges detected. + """ + + # Denoise the image: + processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + processed_image = cv2.GaussianBlur( + processed_image, + ksize = (5, 5), + sigmaX = 0 + ) + + # Detect edges in the image: + processed_image = cv2.Canny( + processed_image, + threshold1 = 50, + threshold2 = 200 + ) + + # Close small gaps in the edges: + kernel = np.ones((5, 5), np.uint8) + processed_image = cv2.dilate(processed_image, kernel, iterations = 1) + # processed_image = cv2.morphologyEx(processed_image, cv2.MORPH_CLOSE, kernel) + # processed_image = cv2.erode(processed_image, kernel, iterations = 1) + + # Done here: + return processed_image + + def scan(self, image, min_area = 0.125): + + """ + Looks for a rectangle in an image and flattens it out. No colour filters are applied here. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :param min_area: The minimum area occupied by the document in the image. + :return: The flattened image if any, else None. + """ + + # Open the image and read the data: + image = self.open_image(image) + if image is None: return None + height, width = image.shape[:2] + image_area = height * width + + # Enhance the image to have better visibility of edges: + # processed_image = self.enhance(image.copy()) + + # Detect edges in the image: + processed_image = self.get_edges(image.copy()) + + # Find the contours in the edges, + # and sort them in ascending order: + contours = cv2.findContours(processed_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + contours = imutils.grab_contours(contours) + contours = sorted(contours, key = cv2.contourArea, reverse = True) + + # We try to get clean and passable contours. + # For our purposes, 'clean' is when it can be approximated to exactly 4 sides, + # and 'passable' is when the approximation comes close to it: + clean_rects = [] + passable_rects = [] + for contour in contours: + perimeter = cv2.arcLength(contour, closed = True) + approximation = cv2.approxPolyDP(contour, 0.025 * perimeter, closed = True) + if len(approximation) == 4: clean_rects.append(approximation) + elif len(approximation) <= 32: + candidate_contour = cv2.minAreaRect(approximation) + candidate_contour = cv2.boxPoints(candidate_contour) + candidate_contour = np.intp(candidate_contour) + passable_rects.append(candidate_contour) + clean_rects = sorted(clean_rects, key = cv2.contourArea, reverse = True) + passable_rects = sorted(passable_rects, key = cv2.contourArea, reverse = True) + + # Now we decide between the best candidate: + document_outline = None + if len(clean_rects) > 0: + candidate_rect = np.array([p[0] for p in clean_rects[0]]) + rectangle_area = cv2.contourArea(candidate_rect) + if rectangle_area / image_area >= min_area: document_outline = candidate_rect + if document_outline is None and len(passable_rects) > 0: + candidate_rect = cv2.minAreaRect(passable_rects[0]) + candidate_rect = cv2.boxPoints(candidate_rect) + candidate_rect = np.intp(candidate_rect) + rectangle_area = cv2.contourArea(candidate_rect) + if rectangle_area / image_area >= min_area: document_outline = candidate_rect + + # If there is no rectangular contour, + # we exit with failure: + if document_outline is None: return None + + # Since we did get the best candidate for the document, + # we figure out which point n the contour is which corner of the document: + document_rectangle = np.zeros((4, 2), dtype = "float32") + sum_points = document_outline.sum(axis = 1) + document_rectangle[0] = document_outline[np.argmin(sum_points)] + document_rectangle[2] = document_outline[np.argmax(sum_points)] + diff_points = np.diff(document_outline, axis = 1) + document_rectangle[1] = document_outline[np.argmin(diff_points)] + document_rectangle[3] = document_outline[np.argmax(diff_points)] + (top_left, top_right, bottom_right, bottom_left) = document_rectangle + + # let's compute the dimensions of the document in the image: + left_height = np.sqrt(((top_left[0] - bottom_left[0]) ** 2) + ((top_left[1] - bottom_left[1]) ** 2)) + right_height = np.sqrt(((top_right[0] - bottom_right[0]) ** 2) + ((top_right[1] - bottom_right[1]) ** 2)) + top_width = np.sqrt(((top_right[0] - top_left[0]) ** 2) + ((top_right[1] - top_left[1]) ** 2)) + bottom_width = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + ((bottom_right[1] - bottom_left[1]) ** 2)) + max_height = max(int(left_height), int(right_height)) + max_width = max(int(top_width), int(bottom_width)) + + # We compute the destination of the transform: + destination = np.array( + object = [ + [0, 0], + [max_width - 1, 0], + [max_width - 1, max_height - 1], + [0, max_height - 1] + ], + dtype = "float32" + ) + + # We apply the transform that flattens out the document: + matrix = cv2.getPerspectiveTransform(document_rectangle, destination) + flattened_image = cv2.warpPerspective(image, matrix, dsize = (max_width, max_height)) + + # Done here: + return flattened_image + + def extract_text( + self, + image, + min_confidence = 0.5, + margin = 0.05 + ): + + """ + To run OCR on the input image. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :param min_confidence: The minimum amount of confidence for detected text to be considered. + :param margin: The extra pixels to include when cropping into a section for OCR. + :return: The extracted text. + """ + + # Make a variable that will hold the final result: + extracted_text = {} + + # Open the image and read the data: + image = self.open_image(image) + height, width = image.shape[:2] + + # Run OCR on the entire canvas: + ocr_result = [r[1] for r in self.__ocr_engine.readtext(image) if r[-1] >= min_confidence] + extracted_text["fullDoc"] = " ".join(ocr_result) + + # Get the doc layout in the image: + document_sections = self.__yolo.predict(image) + document_sections = document_sections["boxes"] + + # For each section, we run the OCR process individually: + section_results = [] + for section in document_sections: + if section["class"] in self.__whitelisted_yolo_classes: + section_width = section["x2"] - section["x1"] + section_height = section["y2"] - section["y1"] + x1 = max(int(section["x1"] - (margin * section_width)), 0) + y1 = max(int(section["y1"] - (margin * section_height)), 0) + x2 = min(int(section["x2"] + (margin * section_width)), width) + y2 = min(int(section["y2"] + (margin * section_height)), height) + sub_image = image[y1:y2, x1:x2] + ocr_result = [r[1] for r in self.__ocr_engine.readtext(sub_image) if r[-1] >= min_confidence] + section_results.append({ + "text": " ".join(ocr_result), + "class": section["class"], + "className": section["className"], + "x1": x1, + "y1": y1, + "x2": x2, + "y2": y2 + }) + extracted_text["bySection"] = section_results + + # Done here: + return extracted_text + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import time + + my_scanner = DocumentScanner( + layout_detection_yolo = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt", + whitelisted_yolo_classes = [0, 1, 3, 4, 5, 7, 9, 10], + # whitelisted_yolo_classes = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10], + ocr_languages = ["en"] + ) + + image = my_scanner.open_image(r"/home/developer/Downloads/talkaholics_card.jpg") + document_image = my_scanner.scan(image) + if document_image is not None: + cv2.imwrite(r"/home/developer/Downloads/flattened_image.jpg", document_image) + start_time = time.time() + document_text = my_scanner.extract_text( + document_image, + min_confidence = 0.5 + ) + print("OCR RESULT:") + print(document_text) + + print(f"FINISHED IN {time.time() - start_time} SECONDS!") + + else: print("No image") diff --git a/utils_v2/image/scanner/scanner_v2.py b/utils_v2/image/scanner/scanner_v2.py new file mode 100644 index 0000000..fc63627 --- /dev/null +++ b/utils_v2/image/scanner/scanner_v2.py @@ -0,0 +1,379 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Tuesday, 24th Sept., 2024 + + OBJECTIVE: + + To provide a way to easily scan documents like printouts and visiting cards. + + REFERENCES: + + 01. https://medium.com/@victorolufemi/build-a-document-scanner-with-opencv-ff9f645a4085 + 02. https://github.com/JaidedAI/ + 03. https://github.com/criistian14/flutter_document_scanner/blob/master/flutter_document_scanner_android/android/src/main/kotlin/com/christian/flutterDocumentScanner/OpenCVPlugin.kt + 04. https://www.geeksforgeeks.org/python-bilateral-filtering/ + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# System-level activities: +import io + +# To work with image processing: +import cv2 +import imutils +import numpy as np + +# To run OCR: +import easyocr + +# To download images from the web: +import requests + +# My utils: +from utils_v2.string import json +from utils_v2.ai.object_detection.yolo import YoloDetect + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class DocumentScanner: + + def __init__( + self, + layout_detection_yolo = None, + whitelisted_yolo_classes = None, + ocr_languages = None + ): + + self.__yolo = YoloDetect( + model_file = layout_detection_yolo, + debug = False + ) + self.__whitelisted_yolo_classes = whitelisted_yolo_classes or [] + self.__ocr_engine = easyocr.Reader(ocr_languages or ["en"]) + pass + + @staticmethod + def open_image(image): + + """ + Opens an image from various sources. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :return: The image as a CV2 object. + """ + + # When the image is provided as a string, + # it could be either as a URL or a path to a local file: + if isinstance(image, str): + if image.startswith("https://") or image.startswith("http://"): + image = io.BytesIO(requests.get(image).content) + else: image = cv2.imread(image) + + # If the image is provided as a io.BytesIO buffer: + if isinstance(image, io.BytesIO): + image.seek(0) + image = np.asarray(bytearray(image.read()), dtype = np.uint8) + image = cv2.imdecode(image, cv2.IMREAD_COLOR) + + # Done here: + return image + + @staticmethod + def show(image, title = "Preview", wait = True): + + """ + Just a quick wrapper to show the image in a window. + :param image: The image that you want to show. + :param title: The title of the window. + :param wait: Set this to True when you want to show the window(s). This is useful when you want to show many + windows at once. Suppose you want to show a lot of windows, you set this to False for all calls except the + very last one. + :return: None. + """ + + cv2.imshow(title, image) + if wait: + cv2.waitKey() + cv2.destroyAllWindows() + + @staticmethod + def get_edges(image): + + """ + Apply all the preprocessing filters on the image before sending it off for contour-finding. + :param image: The image to pre-process. + :return: The pre-processed image. + """ + + # Denoise the image: + processed_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + processed_image = cv2.bilateralFilter( + processed_image, + d = 9, + sigmaColor = 75, + sigmaSpace = 75 + ) + + # Detect edges in the image: + processed_image = cv2.Canny( + processed_image, + threshold1 = 75, + threshold2 = 200 + ) + + # Close small gaps in the edges: + kernel = np.ones((5, 5), np.uint8) + processed_image = cv2.dilate(processed_image, kernel, iterations = 1) + # processed_image = cv2.morphologyEx(processed_image, cv2.MORPH_CLOSE, kernel) + # processed_image = cv2.erode(processed_image, kernel, iterations = 1) + + # Done here: + return processed_image + + def scan(self, image, min_area = 0.125): + + """ + Looks for a rectangle in an image and flattens it out. No colour filters are applied here. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :param min_area: The minimum area occupied by the document in the image. + :return: The flattened image if any, else None. + """ + + # Open the image and read the data: + image = self.open_image(image) + if image is None: return None + height, width = image.shape[:2] + image_area = height * width + + # image = self.enhance(image) + processed_image = self.get_edges(image) + # self.show(processed_image) + + # Find the contours in the edges, + # and sort them in ascending order: + contours = cv2.findContours(processed_image, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + contours = imutils.grab_contours(contours) + contours = sorted(contours, key = cv2.contourArea, reverse = True) + + # We try to get clean and passable contours. + # For our purposes, 'clean' is when it can be approximated to exactly 4 sides, + # and 'passable' is when the approximation comes close to it: + clean_rects = [] + passable_rects = [] + for contour in contours: + perimeter = cv2.arcLength(contour, closed = True) + approximation = cv2.approxPolyDP(contour, 0.025 * perimeter, closed = True) + if len(approximation) == 4: clean_rects.append(approximation) + elif len(approximation) <= 32: + candidate_contour = cv2.minAreaRect(approximation) + candidate_contour = cv2.boxPoints(candidate_contour) + candidate_contour = np.intp(candidate_contour) + passable_rects.append(candidate_contour) + clean_rects = sorted(clean_rects, key = cv2.contourArea, reverse = True) + passable_rects = sorted(passable_rects, key = cv2.contourArea, reverse = True) + + # Now we decide between the best candidate: + document_outline = None + if len(clean_rects) > 0: + candidate_rect = np.array([p[0] for p in clean_rects[0]]) + rectangle_area = cv2.contourArea(candidate_rect) + if rectangle_area / image_area >= min_area: document_outline = candidate_rect + if document_outline is None and len(passable_rects) > 0: + candidate_rect = cv2.minAreaRect(passable_rects[0]) + candidate_rect = cv2.boxPoints(candidate_rect) + candidate_rect = np.intp(candidate_rect) + rectangle_area = cv2.contourArea(candidate_rect) + if rectangle_area / image_area >= min_area: document_outline = candidate_rect + + # If there is no rectangular contour, + # we exit with failure: + if document_outline is None: return None + + # Since we did get the best candidate for the document, + # we figure out which point n the contour is which corner of the document: + document_rectangle = np.zeros((4, 2), dtype = "float32") + sum_points = document_outline.sum(axis = 1) + document_rectangle[0] = document_outline[np.argmin(sum_points)] + document_rectangle[2] = document_outline[np.argmax(sum_points)] + diff_points = np.diff(document_outline, axis = 1) + document_rectangle[1] = document_outline[np.argmin(diff_points)] + document_rectangle[3] = document_outline[np.argmax(diff_points)] + (top_left, top_right, bottom_right, bottom_left) = document_rectangle + + # let's compute the dimensions of the document in the image: + left_height = np.sqrt(((top_left[0] - bottom_left[0]) ** 2) + ((top_left[1] - bottom_left[1]) ** 2)) + right_height = np.sqrt(((top_right[0] - bottom_right[0]) ** 2) + ((top_right[1] - bottom_right[1]) ** 2)) + top_width = np.sqrt(((top_right[0] - top_left[0]) ** 2) + ((top_right[1] - top_left[1]) ** 2)) + bottom_width = np.sqrt(((bottom_right[0] - bottom_left[0]) ** 2) + ((bottom_right[1] - bottom_left[1]) ** 2)) + max_height = max(int(left_height), int(right_height)) + max_width = max(int(top_width), int(bottom_width)) + + # We compute the destination of the transform: + destination = np.array( + object = [ + [0, 0], + [max_width - 1, 0], + [max_width - 1, max_height - 1], + [0, max_height - 1] + ], + dtype = "float32" + ) + + # We apply the transform that flattens out the document: + matrix = cv2.getPerspectiveTransform(document_rectangle, destination) + flattened_image = cv2.warpPerspective(image, matrix, dsize = (max_width, max_height)) + + # Done here: + return flattened_image + + def extract_text( + self, + image, + min_confidence = 0.5, + margin = 0.05 + ): + + """ + To run OCR on the input image. + :param image: The image either as a path to a file on the local disk, or a URL, or an io.BytesIO buffer. + :param min_confidence: The minimum amount of confidence for detected text to be considered. + :param margin: The extra pixels to include when cropping into a section for OCR. + :return: The extracted text. + """ + + # Make a variable that will hold the final result: + extracted_text = {} + + # Open the image and read the data: + image = self.open_image(image) + height, width = image.shape[:2] + + # Run OCR on the entire canvas: + ocr_result = [r[1] for r in self.__ocr_engine.readtext(image) if r[-1] >= min_confidence] + extracted_text["fullDoc"] = " ".join(ocr_result) + + # Get the doc layout in the image: + document_sections = self.__yolo.predict(image) + document_sections = document_sections["boxes"] + + # For each section, we run the OCR process individually: + section_results = [] + for section in document_sections: + if section["class"] in self.__whitelisted_yolo_classes: + section_width = section["x2"] - section["x1"] + section_height = section["y2"] - section["y1"] + x1 = max(int(section["x1"] - (margin * section_width)), 0) + y1 = max(int(section["y1"] - (margin * section_height)), 0) + x2 = min(int(section["x2"] + (margin * section_width)), width) + y2 = min(int(section["y2"] + (margin * section_height)), height) + sub_image = image[y1:y2, x1:x2] + ocr_result = [r[1] for r in self.__ocr_engine.readtext(sub_image) if r[-1] >= min_confidence] + section_results.append({ + "text": " ".join(ocr_result), + "class": section["class"], + "className": section["className"], + "x1": x1, + "y1": y1, + "x2": x2, + "y2": y2 + }) + extracted_text["bySection"] = section_results + + # Done here: + return extracted_text + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + import time + + my_scanner = DocumentScanner( + layout_detection_yolo = r"/home/developer/PycharmProjects/utils/data/ai/behaviour_models/hugging_face/object_detection/YOLOv10-Document-Layout-Analysis/yolov10x_best.pt", + whitelisted_yolo_classes = [0, 1, 3, 4, 5, 7, 9, 10], + # whitelisted_yolo_classes = [0, 1, 3, 4, 5, 6, 7, 8, 9, 10], + ocr_languages = ["en"] + ) + + image = my_scanner.open_image(r"/home/developer/Downloads/sushmita_card.jpg") + # image = my_scanner.open_image(r"/home/developer/Downloads/niranjan_card.jpg") + # image = my_scanner.open_image(r"/home/developer/Downloads/niranjan_card_2.jpg") + # image = my_scanner.open_image(r"/home/developer/Downloads/card_square.jpg") + document_image = my_scanner.scan(image) + if document_image is not None: + cv2.imwrite(r"/home/developer/Downloads/flattened_image.jpg", document_image) + start_time = time.time() + document_text = my_scanner.extract_text( + document_image, + min_confidence = 0.5 + ) + print("OCR RESULT:") + print(document_text) + + print(f"FINISHED IN {time.time() - start_time} SECONDS!") + + else: print("No image") diff --git a/utils_v2/logging/__init__.py b/utils_v2/logging/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/logging/__pycache__/__init__.cpython-310.pyc b/utils_v2/logging/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..b695b94 Binary files /dev/null and b/utils_v2/logging/__pycache__/__init__.cpython-310.pyc differ diff --git a/utils_v2/logging/__pycache__/model.cpython-310.pyc b/utils_v2/logging/__pycache__/model.cpython-310.pyc new file mode 100644 index 0000000..d25e344 Binary files /dev/null and b/utils_v2/logging/__pycache__/model.cpython-310.pyc differ diff --git a/utils_v2/logging/context.py b/utils_v2/logging/context.py new file mode 100644 index 0000000..8de91e0 --- /dev/null +++ b/utils_v2/logging/context.py @@ -0,0 +1,469 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Friday, 2nd Aug., 2024 + + OBJECTIVE: + + To provide an easy way to log all system activities by way of managing the context of what is going on. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + + NOTES: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To manage context: +import contextvars +from contextlib import contextmanager, asynccontextmanager + +# To make decorators: +from functools import wraps + +# For system-level activities: +import os + +# My utils: +from utils_v2.string import json +from utils_v2.datetime import datetime +from utils_v2.security import sanitizers +from utils_v2.api.codes import StatusCodes, HttpCodes +from utils_v2.api.log import APILogModel +from utils_v2.api.response import ResponseModel + +# The needed data models: +from utils_v2.logging.model import GeneralLogModel + +# To work with date and time: +import time +import datetime + +# For random strings: +import random + +# For debugging: +import traceback +import string + +# To work with datatypes: +from types import NoneType +import pandas as pd + +# To work with Pydantic objects: +from pydantic import BaseModel + +# For asynchronous activities: +import asyncio + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# Chars to choose from for random strings: +ALPHANUMERIC_CHARS = string.ascii_letters + string.digits + +# To capture system information: +PROCESS_ID = os.getppid() +PARENT_PROCESS_ID = os.getppid() + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** EXCEPTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +def describe_exception(exc): + + """ + Describes the exception in detail. It extracts the type of exception, a brief message, and even the entire + traceback. Useful for debugging in details without the terminal. You could either log the resultant dict or send it + to the dev team over some service like WhatsApp/Telegram. + :param exc: The exception that occurred. + :return: The dict that explains the exception. + """ + + exc_desc = { + "type": type(exc).__name__, + "msg": str(exc), + "tb": [str(exc_tb) for exc_tb in traceback.format_exception(exc, value = exc, tb = exc.__traceback__)] + } + + return exc_desc + + +# ***************************************************************************************************************** +# ***** **** +# *** Classes *** +# ***** **** +# ***************************************************************************************************************** + + +class AsyncMongoLogger: + + def __init__(self, db_conn, collection = "logs"): + + """ + This class uses an instance of 'AsyncMongo' and makes it usable as a logger. + :param db_conn: The instance of 'AsyncMongo' to use. + :param collection: The collection to write the log into. + """ + + self.__db_conn = db_conn + self.__collection = collection + + async def log(self, log_json): + + """ + Log something to the database using the connection provided when the object was made. + :param log_json: The dict to log. + :return: True if logged successfully, else False. + """ + + asyncio.create_task(self.__db_conn.insert_one( + collection = self.__collection, + document = log_json, + raise_exception = False + )) + + +# --------------------------------------------------------------------------------------------------------------------- + + +class AsyncLoggerContext: + + # Create the context-aware variable(s): + logger = contextvars.ContextVar("logger", default = None) + log_chain = contextvars.ContextVar("log_chain", default = None) + + @classmethod + @asynccontextmanager + async def logging_context(cls, logger, log_chain = None): + + """ + This function makes the context manager that makes the value of the log chain available to everything that is + called within the scope of the context. + :param logger: The object which is to be used to write the log. It should have a 'log' method which should take + in a dict as its input. + :param log_chain: The value of the log chain to be made available within the scope. + :return: None. + """ + + # Set the context: + token_logger = cls.logger.set(logger) + token_log_chain = cls.log_chain.set(log_chain) + + # Make the objects available within the context: + try: yield + + # Release the objects when the context is over: + finally: + cls.logger.reset(token_logger) + cls.log_chain.reset(token_log_chain) + + @staticmethod + def generate_log_id(count = 8): + return "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(min(8, count))) + + @classmethod + def get_logger(cls): + return cls.logger.get() + + @classmethod + def get_log_chain(cls): + return cls.log_chain.get() + + @staticmethod + def summarize( + value, + str_limit = 100, + expand: bool | int = False, + sensitive_keys: list[str] = None + ): + + """ + To summarize an input value to capture the essence without hoarding to much data. + :param value: Anything that you want to summarize. + :param str_limit: The max. no. of chars of a string to retain. + :param expand: Set to True for full expansion, False for no expansion, and an integer for a specific level of + expansion. Applicable on iterables and dicts. The smaller this number, the more concise the summary will be, + and vice versa. + :param sensitive_keys: The list of keys (of a dict) to obscure when summarizing. + :return: The summarized version of the input. + """ + + # If the input is a Pydantic class: + if isinstance(value, BaseModel): value = value.model_dump() + + # Check the sensitive keys: + if sensitive_keys is None: sensitive_keys = [] + + # Handle datatypes that you don't want to modify: + if isinstance(value, (int, float, bool, NoneType)): pass + + # When the value is a list or similar iterable: + elif isinstance(value, (list, tuple, set)): + if expand: + if not isinstance(expand, bool): expand -= 1 + value = [AsyncLoggerContext.summarize( + v, + expand = expand, + sensitive_keys = sensitive_keys + ) for v in value] + else: value = f"array of {len(value)} item(s)" + + # If the value is a dict: + elif isinstance(value, dict): + if expand: + if not isinstance(expand, bool): expand -= 1 + value = { + k: AsyncLoggerContext.summarize( + v, + expand = expand, + sensitive_keys = sensitive_keys + ) if k not in sensitive_keys else "********" + for k, v in value.items() + } + else: value = f"object of {len(value.keys())} field(s) [{', '.join(value.keys())}]" + + # When a dataframe is passed: + elif isinstance(value, pd.DataFrame): + cols = value.columns.to_list() + value = f"table with {len(cols)} col(s) [{', '.join(cols)}] and {len(value)} row(s)" + str_limit = 999 + + # If the input is some form of non-standard object: + else: value = str(value) + + # Handle strings: + if isinstance(value, str): + if len(value) > str_limit: value = value[:str_limit] + "..." + + # Done here: + return value + + @classmethod + def log_it( + cls, + api_version: str = None, + project: str = None, + log_type: str = None, + operation: str = None, + log_input: bool | int = True, + log_output: bool | int = True, + sensitive_keys: list = None + ): + + """ + A decorator factor that can be used to log the results of functions automatically. + :param api_version: A string that indicates the version code of the function being decorated. + :param project: A hint about which project is being worked on. + :param log_type: A hint about which module is being worked on. + :param operation: A hint about which action in a particular module is being worked on. + :param log_input: Set to True to capture everything that went into the function, False to capture the least + info, and set it to an integer to capture a certain depth of the input (applicable on iterables and dicts. + :param log_output: The same as 'log_input', but applicable to the response from the function. + :param sensitive_keys: Keys of a dict whose values must be obscured even if that depth is being captured. + :return: A decorator with the configuration. + """ + + def decorator(func): + + @wraps(func) + async def wrapper(*args, **kwargs): + + # Make variables and extract available info.: + exception = None + response = None + request_ts = date_time.get_current_utc_date_time() + start_ts = time.perf_counter() + cpu_start_ts = time.process_time() + + # Execute the function that is being wrapped: + try: response = await func(*args, **kwargs) + except Exception as exc: exception = exc + + # Do the next steps only if within the logging context: + if cls.get_logger() is not None: + + # Create the log: + if not args: args = None + if not kwargs: kwargs = None + func_log = GeneralLogModel( + pid = PROCESS_ID, + ppid = PARENT_PROCESS_ID, + project = project, + log = log_type or func.__name__, + operation = operation or func.__name__, + apiVer = api_version, + logId = cls.generate_log_id(), + logChain = cls.get_log_chain(), + ts = request_ts, + tat = time.perf_counter() - start_ts, + cpuTime = time.process_time() - cpu_start_ts, + func = func.__name__, + args = cls.summarize(args, expand = log_input, sensitive_keys = sensitive_keys), + kwargs = cls.summarize(kwargs, expand = log_input, sensitive_keys = sensitive_keys), + exception = None if exception is None else describe_exception(exception), + response = cls.summarize(response, expand = log_output, sensitive_keys = sensitive_keys), + ).model_dump() + + # Write the log: + await cls.get_logger().log(func_log) + + # Done here: + if exception is not None: raise exception + return response + + return wrapper + + return decorator + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + from utils_v2.database.async_mongo_v2 import AsyncMongo + + @AsyncLoggerContext.log_it( + api_version = "0.0.1", + project = "testProj", + log_type = "work", + operation = "someWork", + log_input = True, + log_output = True, + sensitive_keys = None + ) + async def some_work(*args, **kwargs): + print("SOME WORK:", AsyncLoggerContext.get_log_chain()) + await asyncio.sleep(max(2.0 * random.random(), 1.0)) + total = sum(args) + return total + + @AsyncLoggerContext.log_it( + api_version = "0.0.1", + project = "testProj", + log_type = "work", + operation = "moreWork", + log_input = True, + log_output = True, + sensitive_keys = ["password", "sessionToken"] + ) + async def more_work(*args, **kwargs): + print("MORE WORK:", AsyncLoggerContext.get_log_chain()) + await asyncio.sleep(max(2.0 * random.random(), 1.0)) + return {"success": True, "sessionToken": "1234567890"} + + @AsyncLoggerContext.log_it( + api_version = "0.0.1", + project = "testProj", + log_type = "work", + operation = "moreWork", + log_input = True, + log_output = True, + sensitive_keys = None + ) + async def last_work(*args, **kwargs): + print("LAST WORK:", AsyncLoggerContext.get_log_chain()) + await asyncio.sleep(max(2.0 * random.random(), 1.0)) + + async def main(chain = None): + + # Connect to MongoDB: + mongo = AsyncMongo( + connection_string = r"mongodb://del.ditscentre.in:27017,wtt.ditscentre.in:27017,mum.arh.001.ditscentre.in:27017/admin?tls=true&tlsCAFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_ca.pem&tlsCertificateKeyFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_cert.pem&replicaSet=dits_mongod_rep&readPreference=primary&authMechanism=MONGODB-X509&authSource=%24external", + database_name = "converse", + max_connections = 10, + debug = True + ) + + # Convert the connection to a logger instance that can be injected + # into the context as a dependency: + mongo_logger = AsyncMongoLogger( + db_conn = mongo, + collection = "logs" + ) + + # Initialize the context: + async with AsyncLoggerContext.logging_context( + logger = mongo_logger, + log_chain = chain + ): + + # Run some functions within the context: + await some_work(1, 2, 3, 4, 5) + await more_work(username = "john.doe@domain.com", password = "mySecretPass") + + # Run something outside the context: + await last_work() + + # async def multi_main(): + # tasks = [ + # main(chain = "kPRwXdItb1"), + # main(chain = "456") + # ] + # await asyncio.gather(*tasks) + + asyncio.run(main(chain = "00wGHRFYPY123")) diff --git a/utils_v2/logging/model.py b/utils_v2/logging/model.py new file mode 100644 index 0000000..355d07e --- /dev/null +++ b/utils_v2/logging/model.py @@ -0,0 +1,131 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Friday, 11th Oct., 2024 + + OBJECTIVE: + + To have a structure for the logs maintained for regular function calls. + This is different from the logs maintained for API calls. + + REFERENCES: + + 01. https://chatgpt.com/share/6708b6c9-6ba4-800f-9ea9-00ec35067512 + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# System-level activities: +import distro +import socket +import platform + +# For data-modelling: +from pydantic import BaseModel, Field +from typing import Any, Optional, List, Literal + +# To work with date and time: +import datetime + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# Info for logging that will stay constant during runtime: +SERVER_HOSTNAME = str(socket.gethostname()) +PLATFORM_INFO = platform.uname() +HOST_OS = str(distro.name(True)) +HOST_CPU = f"{PLATFORM_INFO.processor} ({PLATFORM_INFO.machine})" + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class GeneralLogModel(BaseModel): + + # To identify the machine the code is running on. + # DO NOT MODIFY THESE: + hostname: str = SERVER_HOSTNAME + os: str = HOST_OS + cpu: str = HOST_CPU + # Can modify these: + pid: Optional[Any] = None + ppid: Optional[Any] = None + + # To identify the project and actions: + project: Optional[str] = None + log: str + operation: Optional[str] = None + apiVer: Optional[str] = None + logId: Optional[str] = None + logChain: Optional[str] = None + + # Timing metrics: + ts: datetime.datetime + tat: float + cpuTime: float + + # To understand the inputs: + func: str + args: Optional[Any] = None + kwargs: Optional[Any] = None + + # To understand the outputs: + exception: Optional[Any] = None + response: Optional[Any] = None + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** MAIN PROGRAM *** +# ***** **** +# ***************************************************************************************************************** + + +if __name__ == "__main__": + + pass diff --git a/utils_v2/mail/__init__.py b/utils_v2/mail/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils_v2/mail/async_mail.py b/utils_v2/mail/async_mail.py new file mode 100644 index 0000000..88fe550 --- /dev/null +++ b/utils_v2/mail/async_mail.py @@ -0,0 +1,443 @@ +""" + + AUTHOR: + + Khushal P Soonderji + + DATE: + + Tuesday, 16th Jul, 2024 + + OBJECTIVE: + + To be able to send out mails from code. + + REFERENCES: + + N/A + + DOWNLOADS: + + N/A + +""" + + +# ***************************************************************************************************************** +# ***** **** +# *** IMPORT *** +# ***** **** +# ***************************************************************************************************************** + + +# To make sibling directories accessible for imports: +import sys +sys.path.append(".") +sys.path.append("..") + +# For working with mails: +import aiosmtplib +from email.mime.multipart import MIMEMultipart +from email.mime.image import MIMEImage +from email.mime.text import MIMEText +from email.mime.base import MIMEBase +from email import encoders + +# For random strings: +import string +import random + +# For system-level activities: +import os + +# For working with files in RAM: +import io + +# For debugging: +from icecream import IceCreamDebugger + +# To work with datatypes: +from typing import List + + +# ***************************************************************************************************************** +# ***** **** +# *** MACROS / ONE-TIME INIT *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** VARIABLES *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** FUNCTIONS *** +# ***** **** +# ***************************************************************************************************************** + + +# --- Nothing Yet + + +# ***************************************************************************************************************** +# ***** **** +# *** CLASSES *** +# ***** **** +# ***************************************************************************************************************** + + +class MailMessage: + + def __init__( + self, + to_email: str, + subject: str, + cc_emails: List[str] = None, + bcc_emails: List[str] = None + ): + + """ + Create an instance of the message that you would like to send. + :param to_email: The EMail ID of th recipient. + :param subject: The subject of the mail. + :param cc_emails: A list of recipients to add to the CC section. + :param bcc_emails: A list of recipients to add to the BCC section. + """ + + self.message = MIMEMultipart() + self.message["To"] = to_email + self.message["Subject"] = subject + if cc_emails: self.message["CC"] = ",".join(cc_emails) + if bcc_emails: self.message["BCC"] = ",".join(bcc_emails) + + def add_text(self, text): + + """ + Add plain-text to the mail body. + :param text: The text to add to the mail body. + :return: None. + """ + + self.message.attach(MIMEText(text, "plain")) + + def add_html(self, html_text): + + """ + Add HTML text to the mail body. + :param html_text: The HTML text to add to the mail body. + :return: None. + """ + + self.message.attach(MIMEText(html_text, "html")) + + def add_inline_image(self, image_file, content_id = None): + + """ + Add an inline image to the body of the mail. + NOTE: This is NOT the same as sending an image as an attachment. + :param image_file: The image data to attach to the mail body. + :param content_id: Inline images are inserted via HTML bocks. This field identifies the image resource. If not + specified, I will generate a random string. You may write a custom value here if you know what you are + doing. For most use cases, please ignore this field. + :return: None. + """ + + # Read the image as bytes: + image_bytes = None + if type(image_file) is str: + with open(image_file, "rb") as opened_image_file: + image_bytes = opened_image_file.read() + if type(image_file) is io.BytesIO: + image_file.seek(0) + image_bytes = image_file.getvalue() + + # Declare the part to be attached to the multipart message: + if image_bytes is not None: + + # Create the HTML block if the image pointer is blank: + if content_id is None: + content_id = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(8)) + self.add_html(f""" + +
+