streamlink_manager.py

import os
import streamlink
import shutil
import signal
import sys

class StreamlinkManager:
    def __init__(self, config):
        self.config = config

    M3U8_EXTENSIONS = ['m3u8']

    def get_stream_extension(self, stream):
        if hasattr(stream, 'to_manifest_url'):
            url = stream.to_manifest_url()
            file_extension = url.split('?')[0].split('.')[-1]
            file_extension = stream.url.split('?')[0].split('.')[-1]
            if file_extension in self.M3U8_EXTENSIONS:
                return 'ts'
        else:
            return 'mp4'

    def cleanup(self, fd, temp_filename, final_filename, finalize=True):
        """
        Cleanup function to close the file descriptor and move the temporary file
        to its final destination.
        """
        fd.close()
        # Only finalize (rename .part -> final) on natural end
        if finalize and os.path.exists(temp_filename):
            shutil.move(temp_filename, final_filename)

    def run_streamlink(self, user, recorded_filename, stop_event=None, logger=None):
        session = streamlink.Streamlink()
        session.set_option("twitch-disable-hosting", True)
        session.set_option("twitch-disable-ads", True)
        session.set_option("retry-max", 5)
        session.set_option("retry-streams", 60)

        if self.config.oauth_token:
            session.set_option("http-headers", f"Authorization=OAuth {self.config.oauth_token}")
        quality = self.config.quality
        streams = session.streams(f"twitch.tv/{user}")
        if quality not in streams:
            quality = "best"
        stream = streams[quality]
        extension = self.get_stream_extension(stream)
        temp_filename = f"{recorded_filename}.part"
        final_filename = f"{recorded_filename}.{extension}"
        
        # Open the stream
        fd = stream.open()

        # Note: Signal handlers can only be set in the main thread
        # We'll handle cleanup in the finally block instead

        stopped_by_user = False
        try:
            with open(temp_filename, 'wb') as f:
                while True:
                    if stop_event is not None and stop_event.is_set():
                        stopped_by_user = True
                        if logger: logger.info("Stop requested; breaking read loop")
                        break
                    data = fd.read(1024)
                    if not data:
                        break
                    f.write(data)
        finally:
            # Ensure cleanup is called when the try block exits
            self.cleanup(fd, temp_filename, final_filename, finalize=not stopped_by_user)
        return {"stopped_by_user": stopped_by_user, "final": final_filename, "temp": temp_filename}