U
    ‹¾úgŸ  ã                   @   s@   d dl Z d dlZd dlZd dlZddlmZ G dd„ deƒZdS )é    Né   )ÚFFmpegPostProcessorc                	   @   sd   e Zd ZddiZddiZeddi–Zddd	d
dddddœe–Zddd„Zdd„ Zdd„ Z	dd„ Z
dS )ÚSponsorBlockPPZYoutubeZYouTubeZpoi_highlightZ	HighlightÚchapterZChapterZSponsorzIntermission/Intro AnimationzEndcards/CreditszUnpaid/Self PromotionzPreview/RecapzFiller TangentzInteraction ReminderzNon-Music Section)ZsponsorZintroZoutroZ	selfpromoÚpreviewZfillerZinteractionZmusic_offtopicNúhttps://sponsor.ajay.appc                 C   s>   t  | |¡ t|p| j ¡ ƒ| _t d|¡r0|nd| | _d S )Nz	https?://zhttps://)	r   Ú__init__ÚtupleÚ
CATEGORIESÚkeysÚ_categoriesÚreÚmatchÚ_API_URL)ÚselfÚ
downloaderÚ
categoriesÚapi© r   úE/tmp/pip-unpacked-wheel-q5ljy6pj/yt_dlp/postprocessor/sponsorblock.pyr   !   s    zSponsorBlockPP.__init__c                 C   sR   |d }|| j kr*|  d|› ¡ g |fS |  d¡ |  || d¡¡|d< g |fS )NÚextractor_keyz"SponsorBlock is not supported for zFetching SponsorBlock segmentsÚdurationZsponsorblock_chapters)Ú
EXTRACTORSÚ	to_screenÚ_get_sponsor_chaptersÚget)r   ÚinfoZ	extractorr   r   r   Úrun&   s    

zSponsorBlockPP.runc                    sœ   ˆ  |d ˆj|d  ¡}‡ ‡fdd„‰‡fdd„|D ƒ}t|ƒt|ƒkrTˆ d¡ ‡fdd	„‰‡fd
d„|D ƒ}|s‚ˆ d¡ nˆ dt|ƒ› d¡ |S )NÚidr   c                    s¦   | d }|dkrdS |d dkr(d|d< | d ˆj krF|d  d7  < ˆ rbˆ |d  dkrbˆ |d< | d rztˆ | d  ƒnd}|dk p¤|dk o¤||d |d   d	k S )
NÚsegment)r   r   Fr   r   ÚcategoryZvideoDurationé   gš™™™™™©?)ÚPOI_CATEGORIESÚabs)ÚsZ	start_endZdiff)r   r   r   r   Úduration_filter3   s    z=SponsorBlockPP._get_sponsor_chapters.<locals>.duration_filterc                    s   g | ]}ˆ |ƒr|‘qS r   r   ©Ú.0r$   )r%   r   r   Ú
<listcomp>F   s      z8SponsorBlockPP._get_sponsor_chapters.<locals>.<listcomp>zjSome SponsorBlock segments are from a video of different duration, maybe from an old version of this videoc              	      sP   | d | d  \}}}|dkr&| d nˆ j | }||||| d ||||fgdœS )Nr   r    r   ÚdescriptionZ
actionType)Ú
start_timeZend_timer    ÚtitleÚtyper   )r
   )r$   ÚstartÚendÚcatr+   )r   r   r   Ú
to_chapterJ   s    úz8SponsorBlockPP._get_sponsor_chapters.<locals>.to_chapterc                    s   g | ]}ˆ |ƒ‘qS r   r   r&   )r0   r   r   r(   V   s     z<No matching segments were found in the SponsorBlock databasezFound z& segments in the SponsorBlock database)Ú_get_sponsor_segmentsr   ÚlenZreport_warningr   )r   r   r   ÚsegmentsZduration_matchZsponsor_chaptersr   )r   r%   r   r0   r   r   0   s    
z$SponsorBlockPP._get_sponsor_chaptersc              
   C   s†   t  | d¡¡ ¡ }| j› d|d d… › dtj |t 	| j
¡t 	dddg¡dœ¡ }|  |¡pbg D ]}|d	 |krd|d
   S qdg S )NÚasciiz/api/skipSegments/é   ú?ÚskipZpoir   )Úservicer   ZactionTypesZvideoIDr3   )ÚhashlibÚsha256ÚencodeÚ	hexdigestr   ÚurllibÚparseÚ	urlencodeÚjsonÚdumpsr   Z_download_json)r   Zvideo_idr8   Z
video_hashÚurlÚdr   r   r   r1   ]   s    
ý
z$SponsorBlockPP._get_sponsor_segments)Nr   )Ú__name__Ú
__module__Ú__qualname__r   r"   ZNON_SKIPPABLE_CATEGORIESr
   r   r   r   r1   r   r   r   r   r   	   s2    ÿ ÿ þø	÷

-r   )r9   r@   r   Úurllib.parser=   Zffmpegr   r   r   r   r   r   Ú<module>   s
   