U
    gY                     @   sh  d Z ddlZddlmZ ddlmZ ddlmZmZm	Z	m
Z
mZmZmZmZmZmZmZ ddlmZmZ ddlmZ ddlmZ dd	lmZmZ dd
lmZ ddlmZm Z m!Z!m"Z"m#Z#m$Z$ dZ%dZ&ededZ'ee' eee'  dddZ(ee' e)e'dddZ*e'eeeee' e)f f dddZ+G dd dej,Z-G dd dej.Z/G dd dee!e ef Z0dS )z1This module contains the PicklePersistence class.    N)deepcopy)Path)AnyCallableDictOptionalSetTupleTypeTypeVarUnioncastoverload)BotTelegramObject)FilePathInput)warn)BasePersistencePersistenceInput)ContextTypes)BDCDUDCDCDataConversationDictConversationKeyz/a known bot replaced by PTB's PicklePersistencez2an unknown bot replaced by PTB's PicklePersistenceTelegramObj)bound)clsreturnc                 C   s    |   }t|dd |D S )zsGets all subclasses of the specified object, recursively. from
    https://stackoverflow.com/a/3862957/9706202
    c                 S   s   g | ]}t |D ]}|qqS  )_all_subclasses).0csr    r    C/tmp/pip-unpacked-wheel-swnnwir2/telegram/ext/_picklepersistence.py
<listcomp>+   s     
  z#_all_subclasses.<locals>.<listcomp>)__subclasses__setunion)r   Z
subclassesr    r    r%   r!   &   s    r!   )r   kwargsr   c                 C   s   |  | }|| |S )a  
    This method is used for unpickling. The data, which is in the form a dictionary, is
    converted back into a class. Works mostly the same as :meth:`TelegramObject.__setstate__`.
    This function should be kept in place for backwards compatibility even if the pickling logic
    is changed, since `_custom_reduction` places references to this function into the pickled data.
    )__new____setstate__)r   r*   objr    r    r%   _reconstruct_to.   s    

r.   c                 C   s*   | j dd}t|d |d< t| j|ffS )z
    This method is used for pickling. The bot attribute is preserved so _BotPickler().persistent_id
    works as intended.
    T)Zinclude_privateZ
api_kwargs)Z
_get_attrsdictr.   	__class__)r   datar    r    r%   _custom_reduction:   s    r2   c                       sb   e Zd ZdZeeed fddZeee	ee
e ef f dddZeee ddd	Z  ZS )
_BotPickler_botbotargsr*   c                    s   || _ t j|| d S Nr5   super__init__selfr7   r8   r*   r0   r    r%   r<   I   s    z_BotPickler.__init__)r-   r   c                 C   s   t |tstS t|S )z
        This method is used for pickling. The bot attribute is preserved so
        _BotPickler().persistent_id works as intended.
        )
isinstancer   NotImplementedr2   r>   r-   r    r    r%   reducer_overrideM   s    
z_BotPickler.reducer_overridec                 C   s,   || j krtS t|tr(tddd tS dS )zUsed to 'mark' the Bot, so it can be replaced later. See
        https://docs.python.org/3/library/pickle.html#pickle.Pickler.persistent_id for more info
        zHUnknown bot instance found. Will be replaced by `None` during unpickling   )
stacklevelN)r5   _REPLACED_KNOWN_BOTr@   r   r   _REPLACED_UNKNOWN_BOTrB   r    r    r%   persistent_idY   s    

z_BotPickler.persistent_id)__name__
__module____qualname__	__slots__r   r   r<   r   r	   r   r
   r/   rC   objectr   strrH   __classcell__r    r    r?   r%   r3   F   s   r3   c                       s>   e Zd ZdZeeed fddZeee dddZ	  Z
S )_BotUnpicklerr4   r6   c                    s   || _ t j|| d S r9   r:   r=   r?   r    r%   r<   k   s    z_BotUnpickler.__init__)pidr   c                 C   s(   |t kr| jS |tkrdS tddS )zSReplaces the bot with the current bot if known, else it is replaced by :obj:`None`.Nz,Found unknown persistent id when unpickling!)rF   r5   rG   pickleUnpicklingError)r>   rQ   r    r    r%   persistent_loado   s
    z_BotUnpickler.persistent_load)rI   rJ   rK   rL   r   r   r<   rN   r   rT   rO   r    r    r?   r%   rP   h   s   rP   c                       s  e Zd ZdZdZedEdeee e	e	e
dd	d
ZedFdeee e	e	e
eeeeeef  ddd
ZdGeee e	e	e
eeeeeef  d fdd
ZddddZeedddZddddZeeddddZeeef dddZeeef dddZeddd Zee dd!d"Zeed#d$d%Z ee!ee dd&d'd(Z"eedd)d*d+Z#eedd,d-d.Z$edd/d0d1Z%edd/d2d3Z&edd4d5d6Z'edd7d8d9Z(eedd:d;d<Z)eedd=d>d?Z*edd@dAdBZ+dddCdDZ,  Z-S )HPicklePersistenceaq  Using python's builtin :mod:`pickle` for making your bot persistent.

    Attention:
        The interface provided by this class is intended to be accessed exclusively by
        :class:`~telegram.ext.Application`. Calling any of the methods below manually might
        interfere with the integration of persistence into :class:`~telegram.ext.Application`.

    Note:
        This implementation of :class:`BasePersistence` uses the functionality of the pickle module
        to support serialization of bot instances. Specifically any reference to
        :attr:`~BasePersistence.bot` will be replaced by a placeholder before pickling and
        :attr:`~BasePersistence.bot` will be inserted back when loading the data.

    Examples:
        :any:`Persistent Conversation Bot <examples.persistentconversationbot>`

    .. seealso:: :wiki:`Making Your Bot Persistent <Making-your-bot-persistent>`

    .. versionchanged:: 20.0

        * The parameters and attributes ``store_*_data`` were replaced by :attr:`store_data`.
        * The parameter and attribute ``filename`` were replaced by :attr:`filepath`.
        * :attr:`filepath` now also accepts :obj:`pathlib.Path` as argument.

    Args:
        filepath (:obj:`str` | :obj:`pathlib.Path`): The filepath for storing the pickle files.
            When :attr:`single_file` is :obj:`False` this will be used as a prefix.
        store_data (:class:`~telegram.ext.PersistenceInput`, optional): Specifies which kinds of
            data will be saved by this persistence instance. By default, all available kinds of
            data will be saved.
        single_file (:obj:`bool`, optional): When :obj:`False` will store 5 separate files of
            `filename_user_data`, `filename_bot_data`, `filename_chat_data`,
            `filename_callback_data` and `filename_conversations`. Default is :obj:`True`.
        on_flush (:obj:`bool`, optional): When :obj:`True` will only save to file when
            :meth:`flush` is called and keep data in memory until that happens. When
            :obj:`False` will store data on any transaction *and* on call to :meth:`flush`.
            Default is :obj:`False`.
        context_types (:class:`telegram.ext.ContextTypes`, optional): Pass an instance
            of :class:`telegram.ext.ContextTypes` to customize the types used in the
            ``context`` interface. If not passed, the defaults documented in
            :class:`telegram.ext.ContextTypes` will be used.

            .. versionadded:: 13.6
        update_interval (:obj:`int` | :obj:`float`, optional): The
            :class:`~telegram.ext.Application` will update
            the persistence in regular intervals. This parameter specifies the time (in seconds) to
            wait between two consecutive runs of updating the persistence. Defaults to 60 seconds.

            .. versionadded:: 20.0
    Attributes:
        filepath (:obj:`str` | :obj:`pathlib.Path`): The filepath for storing the pickle files.
            When :attr:`single_file` is :obj:`False` this will be used as a prefix.
        store_data (:class:`~telegram.ext.PersistenceInput`): Specifies which kinds of data will
            be saved by this persistence instance.
        single_file (:obj:`bool`): Optional. When :obj:`False` will store 5 separate files of
            `filename_user_data`, `filename_bot_data`, `filename_chat_data`,
            `filename_callback_data` and `filename_conversations`. Default is :obj:`True`.
        on_flush (:obj:`bool`): Optional. When :obj:`True` will only save to file when
            :meth:`flush` is called and keep data in memory until that happens. When
            :obj:`False` will store data on any transaction *and* on call to :meth:`flush`.
            Default is :obj:`False`.
        context_types (:class:`telegram.ext.ContextTypes`): Container for the types used
            in the ``context`` interface.

            .. versionadded:: 13.6
    )	bot_datacallback_data	chat_datacontext_typesconversationsfilepathon_flushsingle_file	user_dataNTF<   zAPicklePersistence[Dict[Any, Any], Dict[Any, Any], Dict[Any, Any]]r>   r[   
store_datar]   r\   update_intervalc                 C   s   d S r9   r    r`   r    r    r%   r<      s    zPicklePersistence.__init__zPicklePersistence[UD, CD, BD]r>   r[   ra   r]   r\   rb   rY   c                 C   s   d S r9   r    rc   r    r    r%   r<      s    	)r[   ra   r]   r\   rb   rY   c                    sf   t  j||d t|| _|| _|| _d | _d | _d | _d | _	d | _
ttttttf |p\t | _d S )N)ra   rb   )r;   r<   r   r[   r]   r\   r^   rX   rV   rW   rZ   r   r   r   r   r   r   rY   rc   r?   r    r%   r<      s    	
 )r   c              
   C   s  zl| j d}t| j| }W 5 Q R X |d | _|d | _|d| j	 | _	|di | _
|d | _W n tk
r   i | _i | _i | _| j	 | _	d | _
Y nt tjk
r } z| j j}td| d|W 5 d }~X Y n8 tk
r } ztd	| j j |W 5 d }~X Y nX d S )
Nrbr^   rX   rV   rW   rZ   File # does not contain valid pickle data Something went wrong unpickling )r[   openrP   r7   loadr^   rX   getrY   rV   rW   rZ   OSErrorrR   rS   name	TypeError	Exception)r>   filer1   excfilenamer    r    r%   _load_singlefile   s&    


"z"PicklePersistence._load_singlefile)r[   r   c              
   C   s   z4| d }t| j| W  5 Q R  W S Q R X W n~ tk
rJ   Y d S  tjk
r } ztd|j d|W 5 d }~X Y n4 t	k
r } ztd|j |W 5 d }~X Y nX d S )Nrd   re   rf   rg   )
rh   rP   r7   ri   rk   rR   rS   rm   rl   rn   )r>   r[   ro   rp   r    r    r%   
_load_file
  s    ($zPicklePersistence._load_filec              	   C   sN   | j | j| j| j| jd}| jd}t| j|t	j
d| W 5 Q R X d S )N)rZ   r^   rX   rV   rW   wbprotocol)rZ   r^   rX   rV   rW   r[   rh   r3   r7   rR   HIGHEST_PROTOCOLdump)r>   r1   ro   r    r    r%   _dump_singlefile  s    z"PicklePersistence._dump_singlefile)r[   r1   r   c              	   C   s2   | d}t| j|tjd| W 5 Q R X d S )Nrt   ru   )rh   r3   r7   rR   rw   rx   )r>   r[   r1   ro   r    r    r%   
_dump_file!  s    zPicklePersistence._dump_filec                    sF   | j rn4| js4| t| j d}|s,i }|| _ n|   t| j S )zReturns the user_data from the pickle file if it exists or an empty :obj:`dict`.

        Returns:
            Dict[:obj:`int`, :obj:`dict`]: The restored user data.
        
_user_data)r^   r]   rs   r   r[   rr   r   r>   r1   r    r    r%   get_user_data%  s    zPicklePersistence.get_user_datac                    sF   | j rn4| js4| t| j d}|s,i }|| _ n|   t| j S )zReturns the chat_data from the pickle file if it exists or an empty :obj:`dict`.

        Returns:
            Dict[:obj:`int`, :obj:`dict`]: The restored chat data.
        
_chat_data)rX   r]   rs   r   r[   rr   r   r|   r    r    r%   get_chat_data6  s    zPicklePersistence.get_chat_datac                    sL   | j rn:| js:| t| j d}|s2| j  }|| _ n|   t| j S )a  Returns the bot_data from the pickle file if it exists or an empty object of type
        :obj:`dict` | :attr:`telegram.ext.ContextTypes.bot_data`.

        Returns:
            :obj:`dict` | :attr:`telegram.ext.ContextTypes.bot_data`: The restored bot data.
        	_bot_data)rV   r]   rs   r   r[   rY   rr   r   r|   r    r    r%   get_bot_dataG  s    
zPicklePersistence.get_bot_datac                    sT   | j rn4| js4| t| j d}|s,d}|| _ n|   | j dkrJdS t| j S )ad  Returns the callback data from the pickle file if it exists or :obj:`None`.

        .. versionadded:: 13.6

        Returns:
            Tuple[List[Tuple[:obj:`str`, :obj:`float`, Dict[:obj:`str`, :class:`object`]]],
            Dict[:obj:`str`, :obj:`str`]] | :obj:`None`: The restored metadata or :obj:`None`,
            if no data was stored.
        _callback_dataN)rW   r]   rs   r   r[   rr   r   r|   r    r    r%   get_callback_dataY  s    

z#PicklePersistence.get_callback_data)rl   r   c                    sR   | j rn8| js8| t| j d}|s0|i i}|| _ n|   | j |i  S )zReturns the conversations from the pickle file if it exists or an empty dict.

        Args:
            name (:obj:`str`): The handlers name.

        Returns:
            :obj:`dict`: The restored conversations for the handler.
        _conversations)rZ   r]   rs   r   r[   rr   rj   copy)r>   rl   r1   r    r    r%   get_conversationsp  s    	z#PicklePersistence.get_conversations)rl   key	new_stater   c                    sj   | j si | _ | j |i ||kr(dS || j | |< | jsf| js^| t| j d| j  n|   dS )aJ  Will update the conversations for the given handler and depending on :attr:`on_flush`
        save the pickle file.

        Args:
            name (:obj:`str`): The handler's name.
            key (:obj:`tuple`): The key the state is changed for.
            new_state (:class:`object`): The new state for the given key.
        Nr   )	rZ   
setdefaultrj   r\   r]   rz   r   r[   ry   )r>   rl   r   r   r    r    r%   update_conversation  s    z%PicklePersistence.update_conversation)user_idr1   r   c                    sb   | j dkri | _ | j ||kr$dS || j |< | js^| jsV| t| j d| j  n|   dS )a  Will update the user_data and depending on :attr:`on_flush` save the pickle file.

        Args:
            user_id (:obj:`int`): The user the data might have been changed for.
            data (:obj:`dict`): The :attr:`telegram.ext.Application.user_data` ``[user_id]``.
        Nr{   )r^   rj   r\   r]   rz   r   r[   ry   )r>   r   r1   r    r    r%   update_user_data  s    

z"PicklePersistence.update_user_data)chat_idr1   r   c                    sb   | j dkri | _ | j ||kr$dS || j |< | js^| jsV| t| j d| j  n|   dS )a  Will update the chat_data and depending on :attr:`on_flush` save the pickle file.

        Args:
            chat_id (:obj:`int`): The chat the data might have been changed for.
            data (:obj:`dict`): The :attr:`telegram.ext.Application.chat_data` ``[chat_id]``.
        Nr~   )rX   rj   r\   r]   rz   r   r[   ry   )r>   r   r1   r    r    r%   update_chat_data  s    

z"PicklePersistence.update_chat_data)r1   r   c                    sH   | j |krdS || _ | jsD| js<| t| j d| j  n|   dS )zWill update the bot_data and depending on :attr:`on_flush` save the pickle file.

        Args:
            data (:obj:`dict` | :attr:`telegram.ext.ContextTypes.bot_data`): The
                :attr:`telegram.ext.Application.bot_data`.
        Nr   )rV   r\   r]   rz   r   r[   ry   r|   r    r    r%   update_bot_data  s    
z!PicklePersistence.update_bot_datac                    sH   | j |krdS || _ | jsD| js<| t| j d| j  n|   dS )a  Will update the callback_data (if changed) and depending on :attr:`on_flush` save the
        pickle file.

        .. versionadded:: 13.6

        Args:
            data (Tuple[List[Tuple[:obj:`str`, :obj:`float`,                 Dict[:obj:`str`, :class:`object`]]], Dict[:obj:`str`, :obj:`str`]]):
                The relevant data to restore :class:`telegram.ext.CallbackDataCache`.
        Nr   )rW   r\   r]   rz   r   r[   ry   r|   r    r    r%   update_callback_data  s    
z&PicklePersistence.update_callback_data)r   r   c                    sP   | j dkrdS | j |d | jsL| jsD| t| j d| j  n|   dS )zWill delete the specified key from the ``chat_data`` and depending on
        :attr:`on_flush` save the pickle file.

        .. versionadded:: 20.0

        Args:
            chat_id (:obj:`int`): The chat id to delete from the persistence.
        Nr~   )rX   popr\   r]   rz   r   r[   ry   )r>   r   r    r    r%   drop_chat_data  s    	
z PicklePersistence.drop_chat_data)r   r   c                    sP   | j dkrdS | j |d | jsL| jsD| t| j d| j  n|   dS )zWill delete the specified key from the ``user_data`` and depending on
        :attr:`on_flush` save the pickle file.

        .. versionadded:: 20.0

        Args:
            user_id (:obj:`int`): The user id to delete from the persistence.
        Nr{   )r^   r   r\   r]   rz   r   r[   ry   )r>   r   r    r    r%   drop_user_data  s    	
z PicklePersistence.drop_user_data)r   r^   r   c                    s   dS )zDoes nothing.

        .. versionadded:: 13.6
        .. seealso:: :meth:`telegram.ext.BasePersistence.refresh_user_data`
        Nr    )r>   r   r^   r    r    r%   refresh_user_data  s    z#PicklePersistence.refresh_user_data)r   rX   r   c                    s   dS )zDoes nothing.

        .. versionadded:: 13.6
        .. seealso:: :meth:`telegram.ext.BasePersistence.refresh_chat_data`
        Nr    )r>   r   rX   r    r    r%   refresh_chat_data  s    z#PicklePersistence.refresh_chat_data)rV   r   c                    s   dS )zDoes nothing.

        .. versionadded:: 13.6
        .. seealso:: :meth:`telegram.ext.BasePersistence.refresh_bot_data`
        Nr    )r>   rV   r    r    r%   refresh_bot_data  s    z"PicklePersistence.refresh_bot_datac                    s   | j r.| js$| js$| js$| js$| jr|   n| jrN| t| j	 d| j | jrn| t| j	 d| j | jr| t| j	 d| j | jr| t| j	 d| j | jr| t| j	 d| j dS )z/Will save all data in memory to pickle file(s).r{   r~   r   r   r   N)
r]   r^   rX   rV   rW   rZ   ry   rz   r   r[   )r>   r    r    r%   flush  s,    
zPicklePersistence.flush)NTFr_   )NTFr_   N)NTFr_   N).rI   rJ   rK   __doc__rL   r   r   r   r   boolfloatr<   r   r   r   r   r   rr   r   rs   ry   rM   rz   r   intr}   r   r   r   r   rN   r   r   r   r   r   r   r   r   r   r   r   r   r   r   rO   r    r    r?   r%   rU   x   s   C    	            rU   )1r   rR   r   r   pathlibr   typingr   r   r   r   r   r	   r
   r   r   r   r   Ztelegramr   r   Ztelegram._utils.typesr   Ztelegram._utils.warningsr   Ztelegram.extr   r   Ztelegram.ext._contexttypesr   Ztelegram.ext._utils.typesr   r   r   r   r   r   rF   rG   r   r!   r/   r.   r2   Picklerr3   	UnpicklerrP   rU   r    r    r    r%   <module>   s&   4 $"