from abc import ABC, abstractmethod

import pytube as pt

class YouTubeAudioExtractor(ABC):

    @abstractmethod
    def extract(self, url: str, save_path: str) -> str:
        pass
    
class PytubeAudioExtractor(YouTubeAudioExtractor):

    def __init__(self,
                 only_audio: bool = True,
                 extension: str = ".mp3") -> None:
        self.only_audio = only_audio
        self.extension = extension

    def extract(self, url: str,
                save_path: str = "yt_audio") -> str:
        yt = pt.YouTube(url)
        stream = yt.streams.filter(only_audio=self.only_audio)[0]
        filename = save_path + self.extension
        stream.download(filename=filename)
        return filename