utils.py

import gzip
import io
import logging
import json

from typing import NoReturn

from rocketpy._encoders import RocketPyEncoder
from starlette.datastructures import Headers, MutableHeaders
from starlette.types import ASGIApp, Message, Receive, Scope, Send

logger = logging.getLogger(__name__)


def rocketpy_encoder(obj):
    """
    Encode a RocketPy object using official RocketPy encoders.

    This function uses RocketPy's official RocketPyEncoder for complete
    object serialization.

    Args:
        obj: RocketPy object (Environment, Motor, Rocket, Flight)

    Returns:
        Dictionary of encoded attributes
    """

    json_str = json.dumps(
        obj,
        cls=RocketPyEncoder,
        include_outputs=True,
        include_function_data=True,
        discretize=True,
        allow_pickle=False,
    )
    return json.loads(json_str)


def collect_simulation_attributes(flight_obj, flight_simulation_class, rocket_simulation_class, motor_simulation_class, environment_simulation_class):
    """
    Collect attributes from various simulation classes and populate them from the flight object.
    
    Args:
        flight_obj: RocketPy Flight object
        flight_simulation_class: FlightSimulation class
        rocket_simulation_class: RocketSimulation class  
        motor_simulation_class: MotorSimulation class
        environment_simulation_class: EnvironmentSimulation class
    
    Returns:
        Dictionary with all collected attributes
    """
    attributes = rocketpy_encoder(flight_obj)
    
    flight_attributes_list = [
        attr for attr in flight_simulation_class.__annotations__.keys() 
        if attr not in ['message', 'rocket', 'env']
    ]

    rocket_attributes_list = [
        attr for attr in rocket_simulation_class.__annotations__.keys() 
        if attr not in ['message', 'motor']
    ]

    motor_attributes_list = [
        attr for attr in motor_simulation_class.__annotations__.keys() 
        if attr not in ['message']
    ]

    environment_attributes_list = [
        attr for attr in environment_simulation_class.__annotations__.keys()
        if attr not in ['message']
    ]

    try:
        for key in flight_attributes_list:
            try:
                value = getattr(flight_obj, key)
                attributes[key] = value
            except AttributeError as e:
                logger.warning(f"Flight attribute '{key}' not found: {e}")
            except Exception as e:
                logger.error(f"Error getting flight attribute '{key}': {type(e).__name__}: {e}")
    except Exception as e:
        logger.error(f"Error processing flight attributes: {type(e).__name__}: {e}")

    try:
        for key in rocket_attributes_list:
            try:
                value = getattr(flight_obj.rocket, key)
                attributes["rocket"][key] = value
            except AttributeError as e:
                logger.warning(f"Rocket attribute '{key}' not found: {e}")
            except Exception as e:
                logger.error(f"Error getting rocket attribute '{key}': {type(e).__name__}: {e}")
    except Exception as e:
        logger.error(f"Error processing rocket attributes: {type(e).__name__}: {e}")

    try:
        for key in motor_attributes_list:
            try:
                value = getattr(flight_obj.rocket.motor, key)
                attributes["rocket"]["motor"][key] = value
            except AttributeError as e:
                logger.warning(f"Motor attribute '{key}' not found: {e}")
            except Exception as e:
                logger.error(f"Error getting motor attribute '{key}': {type(e).__name__}: {e}")
    except Exception as e:
        logger.error(f"Error processing motor attributes: {type(e).__name__}: {e}")
    
    try:
        for key in environment_attributes_list:
            try:
                value = getattr(flight_obj.env, key)
                attributes["env"][key] = value
            except AttributeError as e:
                logger.warning(f"Environment attribute '{key}' not found: {e}")
            except Exception as e:
                logger.error(f"Error getting environment attribute '{key}': {type(e).__name__}: {e}")
    except Exception as e:
        logger.error(f"Error processing environment attributes: {type(e).__name__}: {e}")

    return rocketpy_encoder(attributes)


class RocketPyGZipMiddleware:
    def __init__(
        self, app: ASGIApp, minimum_size: int = 500, compresslevel: int = 9
    ) -> None:
        self.app = app
        self.minimum_size = minimum_size
        self.compresslevel = compresslevel

    async def __call__(
        self, scope: Scope, receive: Receive, send: Send
    ) -> None:
        if scope["type"] == "http":
            headers = Headers(scope=scope)
            if "gzip" in headers.get("Accept-Encoding", ""):
                responder = GZipResponder(
                    self.app,
                    self.minimum_size,
                    compresslevel=self.compresslevel,
                )
                await responder(scope, receive, send)
                return
        await self.app(scope, receive, send)


class GZipResponder:
    # fork of https://github.com/encode/starlette/blob/master/starlette/middleware/gzip.py
    def __init__(
        self, app: ASGIApp, minimum_size: int, compresslevel: int = 9
    ) -> None:
        self.app = app
        self.minimum_size = minimum_size
        self.send: Send = unattached_send
        self.initial_message: Message = {}
        self.started = False
        self.content_encoding_set = False
        self.gzip_buffer = io.BytesIO()
        self.gzip_file = gzip.GzipFile(
            mode="wb", fileobj=self.gzip_buffer, compresslevel=compresslevel
        )

    async def __call__(
        self, scope: Scope, receive: Receive, send: Send
    ) -> None:
        self.send = send
        with self.gzip_buffer, self.gzip_file:
            await self.app(scope, receive, self.send_with_gzip)

    async def send_with_gzip(self, message: Message) -> None:
        message_type = message["type"]
        if message_type == "http.response.start":
            # Don't send the initial message until we've determined how to
            # modify the outgoing headers correctly.
            self.initial_message = message
            headers = Headers(raw=self.initial_message["headers"])
            self.content_encoding_set = "content-encoding" in headers
        elif (
            message_type == "http.response.body" and self.content_encoding_set
        ):
            if not self.started:
                self.started = True
                await self.send(self.initial_message)
            await self.send(message)
        elif message_type == "http.response.body" and not self.started:
            self.started = True
            body = message.get("body", b"")
            more_body = message.get("more_body", False)
            if ((len(body) < self.minimum_size) and not more_body) or any(
                value == b'application/octet-stream'
                for header, value in self.initial_message["headers"]
            ):
                # Don't apply GZip to small outgoing responses or octet-streams.
                await self.send(self.initial_message)
                await self.send(message)  # pylint: disable=unreachable
            elif not more_body:
                # Standard GZip response.
                self.gzip_file.write(body)
                self.gzip_file.close()
                body = self.gzip_buffer.getvalue()

                headers = MutableHeaders(raw=self.initial_message["headers"])
                headers["Content-Encoding"] = "gzip"
                headers["Content-Length"] = str(len(body))
                headers.add_vary_header("Accept-Encoding")
                message["body"] = body

                await self.send(self.initial_message)
                await self.send(message)  # pylint: disable=unreachable
            else:
                # Initial body in streaming GZip response.
                headers = MutableHeaders(raw=self.initial_message["headers"])
                headers["Content-Encoding"] = "gzip"
                headers.add_vary_header("Accept-Encoding")
                del headers["Content-Length"]

                self.gzip_file.write(body)
                message["body"] = self.gzip_buffer.getvalue()
                self.gzip_buffer.seek(0)
                self.gzip_buffer.truncate()

                await self.send(self.initial_message)
                await self.send(message)  # pylint: disable=unreachable

        elif message_type == "http.response.body":
            # Remaining body in streaming GZip response.
            body = message.get("body", b"")
            more_body = message.get("more_body", False)

            self.gzip_file.write(body)
            if not more_body:
                self.gzip_file.close()

            message["body"] = self.gzip_buffer.getvalue()
            self.gzip_buffer.seek(0)
            self.gzip_buffer.truncate()

            await self.send(message)


async def unattached_send(message: Message) -> NoReturn:
    raise RuntimeError("send awaitable not set")  # pragma: no cover