U
    
ڲg                     @   sX  d Z ddlZddlZddlZddlZddlmZ ddlZddlZddl	Z	ddl	m
Z
 ddl	mZ ddl	mZ zddl	mZ dZW n ek
r   d	ZY nX zddlZdZW n ek
r   d	ZY nX d
ZdZedZedZedZedZedZedZedZdd ZdNddZdOddZdPddZ dQddZ!dRddZ"dSddZ#dTd d!Z$G d"d# d#Z%G d$d% d%Z&G d&d' d'e&Z'G d(d) d)Z(G d*d+ d+Z)G d,d- d-e(Z*G d.d/ d/e(Z+G d0d1 d1Z,G d2d3 d3Z-d4d5 Z.dUd6d7Z/d8d9 Z0d:d; Z1d<d= Z2d>d? Z3d@dA Z4dBdC Z5dDdE Z6dFdG Z7dHdI Z8dJdK Z9G dLdM dMZ:dS )VzFirebase ML module.

This module contains functions for creating, updating, getting, listing,
deleting, publishing and unpublishing Firebase ML models.
    N)parse)_http_client)_utils)
exceptions)storageTFZ_mld   z^[A-Za-z0-9_-]{1,60}$z^[A-Za-z0-9_-]{1,32}$z;^gs://(?P<bucket_name>[a-z0-9_.-]{3,63})/(?P<blob_name>.+)$zm^projects/(?P<project_id>[a-z0-9-]{6,30})/locations/(?P<location_id>[^/]+)/models/(?P<model_id>[A-Za-z0-9]+)$zS^projects/(?P<project_id>[a-z0-9-]{6,30})/models/(?P<model_id>[A-Za-z0-9_-]{1,60})$z;^projects/(?P<project_id>[a-z0-9-]{6,30})/operations/[^/]+$c                 C   s   t | ttS )a   Returns an _MLService instance for an App.

    Args:
      app: A Firebase App instance (or None to use the default App).

    Returns:
      _MLService: An _MLService for the specified App instance.

    Raises:
      ValueError: If the app argument is invalid.
    )r   Zget_app_service_ML_ATTRIBUTE
_MLServiceapp r   5/tmp/pip-unpacked-wheel-p0r7i5ii/firebase_admin/ml.py_get_ml_serviceB   s    r   c                 C   s   t |}tj|| |dS )zCreates a model in the current Firebase project.

    Args:
        model: An ml.Model to create.
        app: A Firebase app instance (or None to use the default app).

    Returns:
        Model: The model that was created in Firebase ML.
    r
   )r   Model	from_dictcreate_modelmodelr   
ml_servicer   r   r   r   Q   s    
r   c                 C   s   t |}tj|| |dS )zUpdates a model's metadata or model file.

    Args:
        model: The ml.Model to update.
        app: A Firebase app instance (or None to use the default app).

    Returns:
        Model: The updated model.
    r
   )r   r   r   update_modelr   r   r   r   r   _   s    
r   c                 C   s    t |}tj|j| dd|dS )a  Publishes a Firebase ML model.

    A published model can be downloaded to client apps.

    Args:
        model_id: The id of the model to publish.
        app: A Firebase app instance (or None to use the default app).

    Returns:
        Model: The published model.
    Tpublishr
   r   r   r   set_publishedmodel_idr   r   r   r   r   publish_modelm   s    r   c                 C   s    t |}tj|j| dd|dS )zUnpublishes a Firebase ML model.

    Args:
        model_id: The id of the model to unpublish.
        app: A Firebase app instance (or None to use the default app).

    Returns:
        Model: The unpublished model.
    Fr   r
   r   r   r   r   r   unpublish_model}   s    
r   c                 C   s   t |}tj|| |dS )zGets the model specified by the given ID.

    Args:
        model_id: The id of the model to get.
        app: A Firebase app instance (or None to use the default app).

    Returns:
     Model: The requested model.
    r
   )r   r   r   	get_modelr   r   r   r   r      s    
r   c                 C   s   t |}t|j| |||dS )a0  Lists the current project's models.

    Args:
        list_filter: a list filter string such as ``tags:'tag_1'``. None will return all models.
        page_size: A number between 1 and 100 inclusive that specifies the maximum
            number of models to return per page. None for default.
        page_token: A next page token returned from a previous page of results. None
            for first page of results.
        app: A Firebase app instance (or None to use the default app).

    Returns:
        ListModelsPage: A (filtered) list of models.
    r
   )r   ListModelsPagelist_models)list_filter	page_size
page_tokenr   r   r   r   r   r       s        r    c                 C   s   t |}||  dS )zDeletes a model from the current project.

    Args:
        model_id: The id of the model you wish to delete.
        app: A Firebase app instance (or None to use the default app).
    N)r   delete_modelr   r   r   r   r$      s    r$   c                   @   s  e Zd ZdZd-ddZed.ddZdd Zd	d
 Zdd Z	e
dd Ze
dd Zejdd Zedd Ze
dd Ze
dd Ze
dd Ze
dd Ze
dd Ze
dd Ze
d d! Zejd"d! Ze
d#d$ Zd/d%d&Ze
d'd( Zejd)d( Zd0d+d,ZdS )1r   aI  A Firebase ML Model object.

    Args:
        display_name: The display name of your model - used to identify your model in code.
        tags: Optional list of strings associated with your model. Can be used in list queries.
        model_format: A subclass of ModelFormat. (e.g. TFLiteFormat) Specifies the model details.
    Nc                 C   s@   d | _ i | _d | _|d k	r || _|d k	r.|| _|d k	r<|| _d S N)_app_data_model_formatdisplay_nametagsmodel_format)selfr)   r*   r+   r   r   r   __init__   s    zModel.__init__c                 C   sL   t |}d}|dd}|dd |r2t|}t|d}||_||_|S )-Create an instance of the object from a dict.NtfliteModel@typer+   )dictpopTFLiteFormatr   r   r'   r&   )clsdatar   	data_copytflite_formatZtflite_format_datar   r   r   r   r      s    

zModel.from_dictc                 C   s   t |}|j| _|j| _d S r%   )r   r   r+   r'   )r,   r6   copyr   r   r   _update_from_dict   s    
zModel._update_from_dictc                 C   s(   t || jr$| j|jko"| j|jkS dS NF)
isinstance	__class__r'   r(   r,   otherr   r   r   __eq__   s    zModel.__eq__c                 C   s   |  | S r%   r@   r>   r   r   r   __ne__   s    zModel.__ne__c                 C   s(   | j dsdS t| j d\}}|S )z&The model's ID, unique to the project.nameN)r'   get_validate_and_parse_name)r,   _r   r   r   r   r      s    zModel.model_idc                 C   s   | j dS )zaThe model's display name, used to refer to the model in code and in
        the Firebase console.displayNamer'   rD   r,   r   r   r   r)      s    zModel.display_namec                 C   s   t || jd< | S )NrG   )_validate_display_namer'   )r,   r)   r   r   r   r)      s    c                 C   s>   | sd S d}t j d}t j | |}t||  d }|S )Nz%Y-%m-%dT%H:%M:%S.%fZr   i  )datetimeutcfromtimestampstrptimeinttotal_seconds)Zdate_stringZ
format_strepochZdatetime_objectZmillisr   r   r   _convert_to_millis   s    zModel._convert_to_millisc                 C   s   t | jddS )zThe time the model was created.Z
createTimeNr   rQ   r'   rD   rI   r   r   r   create_time  s    zModel.create_timec                 C   s   t | jddS )z$The time the model was last updated.Z
updateTimeNrR   rI   r   r   r   update_time	  s    zModel.update_timec                 C   s   | j di di dS )zValidation error message.stateZvalidationErrormessagerH   rI   r   r   r   validation_error  s    zModel.validation_errorc                 C   s   t | jdi dS )zMTrue if the model is published and available for clients to
        download.rU   	published)boolr'   rD   rI   r   r   r   rX     s    zModel.publishedc                 C   s   | j dS )z,The entity tag (ETag) of the model resource.etagrH   rI   r   r   r   rZ     s    z
Model.etagc                 C   s   | j dS )z SHA256 hash of the model binary.Z	modelHashrH   rI   r   r   r   
model_hash  s    zModel.model_hashc                 C   s   | j dS )z.Tag strings, used for filtering query results.r*   rH   rI   r   r   r   r*   #  s    z
Model.tagsc                 C   s   t || jd< | S )Nr*   )_validate_tagsr'   )r,   r*   r   r   r   r*   (  s    c                 C   s$   t | jdo t| jddkS )z:True if the Model object is locked by an active operation.activeOperationsr   )rY   r'   rD   lenrI   r   r   r   locked-  s    zModel.lockedc                 C   sN   | j s
dS t| j}| jdd d}|j||d|d}| | dS )aQ  Waits for the model to be unlocked. (All active operations complete)

        Args:
            max_time_seconds: The maximum number of seconds to wait for the model to unlock.
                (None for no limit)

        Raises:
            exceptions.DeadlineExceeded: If max_time_seconds passed and the model is still locked.
        Nr]   r   rC   T)wait_for_operationmax_time_seconds)r_   r   r&   r'   rD   handle_operationget_operationr:   )r,   ra   r   op_nameZ
model_dictr   r   r   wait_for_unlocked3  s    

zModel.wait_for_unlockedc                 C   s   | j S )zeThe model's ``ModelFormat`` object, which represents the model's
        format and storage location.)r(   rI   r   r   r   r+   G  s    zModel.model_formatc                 C   s   |d k	rt | || _| S r%   )_validate_model_formatr(   )r,   r+   r   r   r   r+   M  s    Fc                 C   s(   t | j}| jr$|| jj|d |S )4Returns a serializable representation of the object.
for_upload)r2   r'   r(   updateas_dictr,   ri   r9   r   r   r   rk   T  s    
zModel.as_dict)NNN)N)N)F)__name__
__module____qualname____doc__r-   classmethodr   r:   r@   rB   propertyr   r)   setterstaticmethodrQ   rS   rT   rW   rX   rZ   r[   r*   r_   re   r+   rk   r   r   r   r   r      sN   




	











r   c                   @   s   e Zd ZdZdddZdS )ModelFormatz?Abstract base class representing a Model Format such as TFLite.Fc                 C   s   t dS rg   NNotImplementedErrorr,   ri   r   r   r   rk   ^  s    zModelFormat.as_dictN)Frm   rn   ro   rp   rk   r   r   r   r   ru   \  s   ru   c                   @   sr   e Zd ZdZdddZedd Zdd Zd	d
 Ze	dd Z
edd Zejdd Zedd ZdddZdS )r4   zModel format representing a TFLite model.

    Args:
        model_source: A TFLiteModelSource sub class. Specifies the details of the model source.
    Nc                 C   s   i | _ d | _|d k	r|| _d S r%   )r'   _model_sourcemodel_sourcer,   r|   r   r   r   r-   i  s    zTFLiteFormat.__init__c                 C   s"   t |}t| |d}||_|S )r.   )r|   )r2   r4   _init_model_sourcer'   )r5   r6   r7   r8   r   r   r   r   p  s    zTFLiteFormat.from_dictc                 C   s(   t || jr$| j|jko"| j|jkS dS r;   )r<   r=   r'   r{   r>   r   r   r   r@   x  s    zTFLiteFormat.__eq__c                 C   s   |  | S r%   rA   r>   r   r   r   rB   ~  s    zTFLiteFormat.__ne__c                 C   sD   |  dd}|rt|dS |  dd}|r@tdt t|dS dS )zInitialize the ML model source.gcsTfliteUriN)gcs_tflite_uriautomlModelQAutoML model support is deprecated and will be removed in the next major version.)auto_ml_model)r3   TFLiteGCSModelSourcewarningswarnDeprecationWarningTFLiteAutoMlSource)r6   r   r   r   r   r   r~     s    

zTFLiteFormat._init_model_sourcec                 C   s   | j S )zThe TF Lite model's location.)r{   rI   r   r   r   r|     s    zTFLiteFormat.model_sourcec                 C   s$   |d k	rt |tstd|| _d S )Nz0Model source must be a TFLiteModelSource object.)r<   TFLiteModelSource	TypeErrorr{   r}   r   r   r   r|     s    
c                 C   s   | j dS )z'The size in bytes of the TF Lite model.Z	sizeBytesrH   rI   r   r   r   
size_bytes  s    zTFLiteFormat.size_bytesFc                 C   s,   t | j}| jr$|| jj|d d|iS )rg   rh   r/   )r2   r'   r{   rj   rk   rl   r   r   r   rk     s    
zTFLiteFormat.as_dict)N)F)rm   rn   ro   rp   r-   rq   r   r@   rB   rt   r~   rr   r|   rs   r   rk   r   r   r   r   r4   c  s   





r4   c                   @   s   e Zd ZdZdddZdS )r   zIAbstract base class representing a model source for TFLite format models.Fc                 C   s   t dS rv   rw   ry   r   r   r   rk     s    zTFLiteModelSource.as_dictN)Frz   r   r   r   r   r     s   r   c                   @   sH   e Zd ZdZdZdZedd Zedd Zedd	 Z	ed
d Z
dS )_CloudStorageClientzCloud Storage helper classzgs://{0}/{1}zFirebase/ML/Models/{0}c                   C   s   t stdd S )NznFailed to import the Cloud Storage library for Python. Make sure to install the "google-cloud-storage" module.)_GCS_ENABLEDImportErrorr   r   r   r   _assert_gcs_enabled  s    z'_CloudStorageClient._assert_gcs_enabledc                 C   s*   t | }|std|d|dfS )N!GCS TFLite URI format is invalid.bucket_name	blob_name)_GCS_TFLITE_URI_PATTERNmatch
ValueErrorgroup)urimatcherr   r   r   _parse_gcs_tflite_uri  s    
z)_CloudStorageClient._parse_gcs_tflite_uric                 C   sR   t   tj|}tj| |d}t j|}|	|}|
| t j|j|S )z4Upload a model file to the specified Storage bucket.r
   )r   r   ospathbasenamer   bucket	BLOB_NAMEformatblobZupload_from_filenameGCS_URIrC   )r   model_file_namer   	file_namer   r   r   r   r   r   upload  s    

z_CloudStorageClient.uploadc                 C   sF   t   t | \}}tj||d}||}|jdtjddddS )zHMakes the gcs_tflite_uri readable for GET for 10 minutes via signed_uri.r
   Zv4
   )minutesGET)versionZ
expirationmethod)	r   r   r   r   r   r   Zgenerate_signed_urlrK   	timedelta)r   r   r   r   r   r   r   r   r   sign_uri  s    

z_CloudStorageClient.sign_uriN)rm   rn   ro   rp   r   r   rt   r   r   r   r   r   r   r   r   r     s   


r   c                   @   s   e Zd ZdZe ZdddZdd Zdd Ze	dd	d
Z
edd Zedd Zedd Ze	d ddZe	d!ddZedd Zejdd Zdd Zd"ddZdS )#r   zCTFLite model source representing a tflite model file stored in GCS.Nc                 C   s   || _ t|| _d S r%   )r&   _validate_gcs_tflite_uri_gcs_tflite_uri)r,   r   r   r   r   r   r-     s    zTFLiteGCSModelSource.__init__c                 C   s   t || jr| j|jkS dS r;   )r<   r=   r   r>   r   r   r   r@     s    zTFLiteGCSModelSource.__eq__c                 C   s   |  | S r%   rA   r>   r   r   r   rB     s    zTFLiteGCSModelSource.__ne__c                 C   s   t j|||}t ||dS )a  Uploads the model file to an existing Google Cloud Storage bucket.

        Args:
            model_file_name: The name of the model file.
            bucket_name: The name of an existing bucket. None to use the default bucket configured
                in the app.
            app: A Firebase app instance (or None to use the default app).

        Returns:
            TFLiteGCSModelSource: The source created from the model_file

        Raises:
            ImportError: If the Cloud Storage Library has not been installed.
        )r   r   )r   _STORAGE_CLIENTr   )r5   r   r   r   Zgcs_urir   r   r   from_tflite_model_file  s    z+TFLiteGCSModelSource.from_tflite_model_filec                   C   s>   t stdtjjds:tjjds:tdtjjd S )Nz_Failed to import the tensorflow library for Python. Make sure to install the tensorflow module.1.z2.z5Expected tensorflow version 1.x or 2.x, but found {0})_TF_ENABLEDr   tfr   VERSION
startswithr   r   r   r   r   _assert_tf_enabled  s    z'TFLiteGCSModelSource._assert_tf_enabledc                 C   s   t jj| }| S r%   )r   liteTFLiteConverterfrom_saved_modelconvert)saved_model_dir	converterr   r   r   _tf_convert_from_saved_model
  s    z1TFLiteGCSModelSource._tf_convert_from_saved_modelc                 C   sH   t jjdr2d}t jj| | t jj	|}nt jj
| }| S )z4Converts the given Keras model into a TF Lite model.r   zfirebase_keras_model.h5)r   r   r   r   ZkerasmodelsZ
save_modelr   r   Zfrom_keras_model_filefrom_keras_modelr   )keras_modelZ
keras_filer   r   r   r   _tf_convert_from_keras_model  s    z1TFLiteGCSModelSource._tf_convert_from_keras_modelfirebase_ml_model.tflitec              	   C   s@   t   t |}t|d}|| W 5 Q R X t |||S )a  Creates a Tensor Flow Lite model from the saved model, and uploads the model to GCS.

        Args:
            saved_model_dir: The saved model directory.
            model_file_name: The name that the tflite model will be saved as in Cloud Storage.
            bucket_name: The name of an existing bucket. None to use the default bucket configured
                in the app.
            app: Optional. A Firebase app instance (or None to use the default app)

        Returns:
            TFLiteGCSModelSource: The source created from the saved_model_dir

        Raises:
            ImportError: If the Tensor Flow or Cloud Storage Libraries have not been installed.
        wb)r   r   r   openwriter   )r5   r   r   r   r   tflite_model
model_filer   r   r   r     s
    
z%TFLiteGCSModelSource.from_saved_modelc              	   C   s@   t   t |}t|d}|| W 5 Q R X t |||S )a  Creates a Tensor Flow Lite model from the keras model, and uploads the model to GCS.

        Args:
            keras_model: A tf.keras model.
            model_file_name: The name that the tflite model will be saved as in Cloud Storage.
            bucket_name: The name of an existing bucket. None to use the default bucket configured
                in the app.
            app: Optional. A Firebase app instance (or None to use the default app)

        Returns:
            TFLiteGCSModelSource: The source created from the keras_model

        Raises:
            ImportError: If the Tensor Flow or Cloud Storage Libraries have not been installed.
        r   )r   r   r   r   r   r   )r5   r   r   r   r   r   r   r   r   r   r   5  s
    
z%TFLiteGCSModelSource.from_keras_modelc                 C   s   | j S )z'URI of the model file in Cloud Storage.)r   rI   r   r   r   r   M  s    z#TFLiteGCSModelSource.gcs_tflite_uric                 C   s   t || _d S r%   )r   r   )r,   r   r   r   r   r   R  s    c                 C   s   t j| j| jS )zQSigns the GCS uri, so the model file can be uploaded to Firebase ML and verified.)r   r   r   r   r&   rI   r   r   r   _get_signed_gcs_tflite_uriV  s    z/TFLiteGCSModelSource._get_signed_gcs_tflite_uriFc                 C   s   |rd|   iS d| jiS )rg   r   )r   r   ry   r   r   r   rk   Z  s    zTFLiteGCSModelSource.as_dict)N)NN)r   NN)r   NN)F)rm   rn   ro   rp   r   r   r-   r@   rB   rq   r   rt   r   r   r   r   r   rr   r   rs   r   rk   r   r   r   r   r     s:   



        

r   c                   @   sN   e Zd ZdZdddZdd Zdd Zed	d
 Zej	dd
 ZdddZ
dS )r   zTFLite model source representing a tflite model created with AutoML.

    AutoML model support is deprecated and will be removed in the next major version.
    Nc                 C   s   t dt || _|| _d S )Nr   )r   r   r   r&   r   )r,   r   r   r   r   r   r-   h  s
    zTFLiteAutoMlSource.__init__c                 C   s   t || jr| j|jkS dS r;   )r<   r=   r   r>   r   r   r   r@   n  s    zTFLiteAutoMlSource.__eq__c                 C   s   |  | S r%   rA   r>   r   r   r   rB   s  s    zTFLiteAutoMlSource.__ne__c                 C   s   | j S )zGResource name of the model, created by the AutoML API or Cloud console._auto_ml_modelrI   r   r   r   r   v  s    z TFLiteAutoMlSource.auto_ml_modelc                 C   s   t || _d S r%   )_validate_auto_ml_modelr   )r,   r   r   r   r   r   {  s    Fc                 C   s
   d| j iS )rg   r   r   ry   r   r   r   rk     s    zTFLiteAutoMlSource.as_dict)N)F)rm   rn   ro   rp   r-   r@   rB   rr   r   rs   rk   r   r   r   r   r   b  s   


r   c                   @   sX   e Zd ZdZdd Zedd Zedd Zedd	 Zed
d Z	dd Z
dd ZdS )r   aF  Represents a page of models in a Firebase project.

    Provides methods for traversing the models included in this page, as well as
    retrieving subsequent pages of models. The iterator returned by
    ``iterate_all()`` can be used to iterate through all the models in the
    Firebase project starting from this page.
    c                 C   s0   || _ || _|| _|| _|| _||||| _d S r%   )_list_models_func_list_filter
_page_sizeZ_page_tokenr&   _list_response)r,   Zlist_models_funcr!   r"   r#   r   r   r   r   r-     s    zListModelsPage.__init__c                    s    fdd j dg D S )z A list of Models from this page.c                    s   g | ]}t j| jd qS )r
   )r   r   r&   ).0r   rI   r   r   
<listcomp>  s    z)ListModelsPage.models.<locals>.<listcomp>r   r   rD   rI   r   rI   r   r     s    
zListModelsPage.modelsc                 C   s   | j S )z,The filter string used to filter the models.)r   rI   r   r   r   r!     s    zListModelsPage.list_filterc                 C   s   | j ddS )z+Token identifying the next page of results.ZnextPageToken r   rI   r   r   r   next_page_token  s    zListModelsPage.next_page_tokenc                 C   s
   t | jS )z!True if more pages are available.)rY   r   rI   r   r   r   has_next_page  s    zListModelsPage.has_next_pagec                 C   s$   | j r t| j| j| j| j| jS dS )zRetrieves the next page of models if available.

        Returns:
            ListModelsPage: Next page of models, or None if this is the last page.
        N)r   r   r   r   r   r   r&   rI   r   r   r   get_next_page  s    zListModelsPage.get_next_pagec                 C   s   t | S )aC  Retrieves an iterator for Models.

        Returned iterator will iterate through all the models in the Firebase
        project starting from this page. The iterator will never buffer more than
        one page of models in memory at a time.

        Returns:
            iterator: An iterator of Model instances.
        )_ModelIteratorrI   r   r   r   iterate_all  s    
zListModelsPage.iterate_allN)rm   rn   ro   rp   r-   rr   r   r!   r   r   r   r   r   r   r   r   r     s   



r   c                   @   s0   e Zd ZdZdd Zdd Zdd Zdd	 Zd
S )r   a  An iterator that allows iterating over models, one at a time.

    This implementation loads a page of models into memory, and iterates on them.
    When the whole page has been traversed, it loads another page. This class
    never keeps more than one page of entries in memory.
    c                 C   s"   t |tstd|| _d| _d S )Nz%Current page must be a ListModelsPager   )r<   r   r   _current_page_index)r,   Zcurrent_pager   r   r   r-     s    
z_ModelIterator.__init__c                 C   sf   | j t| jjkr,| jjr,| j | _d| _ | j t| jjk r^| jj| j  }|  j d7  _ |S td S )Nr      )r   r^   r   r   r   r   StopIteration)r,   resultr   r   r   next  s    z_ModelIterator.nextc                 C   s   |   S r%   )r   rI   r   r   r   __next__  s    z_ModelIterator.__next__c                 C   s   | S r%   r   rI   r   r   r   __iter__  s    z_ModelIterator.__iter__N)rm   rn   ro   rp   r-   r   r   r   r   r   r   r   r     s
   r   c                 C   s*   t | }|std|d|dfS )N&Model resource name format is invalid.
project_idr   )_RESOURCE_NAME_PATTERNr   r   r   )rC   r   r   r   r   rE     s    
rE   c                 C   s,   t | tstd|d kr(| js(tdd S )NzModel must be an ml.Model.zModel must have a display name.)r<   r   r   r)   r   )r   update_maskr   r   r   _validate_model  s    
r   c                 C   s   t | stdd S )NzModel ID format is invalid.)_MODEL_ID_PATTERNr   r   )r   r   r   r   _validate_model_id  s    
r   c                 C   s   t | std| S )Nz!Operation name format is invalid.)_OPERATION_NAME_PATTERNr   r   )rd   r   r   r   _validate_operation_name  s    
r   c                 C   s   t | std| S )NzDisplay name format is invalid.)_DISPLAY_NAME_PATTERNr   r   )r)   r   r   r   rJ     s    
rJ   c                 C   sB   t | trtdd | D s$tdtdd | D s>td| S )Nc                 s   s   | ]}t |tV  qd S r%   )r<   strr   tagr   r   r   	<genexpr>
  s     z!_validate_tags.<locals>.<genexpr>zTags must be a list of strings.c                 s   s   | ]}t |V  qd S r%   )_TAG_PATTERNr   r   r   r   r   r     s     zTag format is invalid.)r<   listallr   r   )r*   r   r   r   r\     s    
r\   c                 C   s   t | std| S )Nr   )r   r   r   )r   r   r   r   r     s    
r   c                 C   s   t | std| S )Nr   )_AUTO_ML_MODEL_PATTERNr   r   )r   r   r   r   r     s    
r   c                 C   s   t | tstd| S )Nz*Model format must be a ModelFormat object.)r<   ru   r   r1   r   r   r   rf     s    
rf   c                 C   s   | d k	rt | tstdd S )Nz%List filter must be a string or None.r<   r   r   )r!   r   r   r   _validate_list_filter$  s    
r   c                 C   s>   | d k	r:t | tk	rtd| dk s,| tkr:tdtd S )Nz#Page size must be a number or None.r   z6Page size must be a positive integer between 1 and {0})typerN   r   _MAX_PAGE_SIZEr   r   )r"   r   r   r   _validate_page_size*  s    r   c                 C   s   | d k	rt | tstdd S )Nz$Page token must be a string or None.r   )r#   r   r   r   _validate_page_token4  s    
r   c                   @   st   e Zd ZdZdZdZdZdZdd Zdd	 Z	d
d Z
dddZdd ZdddZdd Zdd Zdd Zdd ZdS )r	   zFirebase ML service.z7https://firebaseml.googleapis.com/v1beta2/projects/{0}/z*https://firebaseml.googleapis.com/v1beta2/g      ?   c                 C   sn   |j | _| jstdtj| j| _ddtji}t	j
|j || jd| _t	j
|j |tjd| _d S )NzqProject ID is required to access ML service. Either set the projectId option, or use service account credentials.zX-FIREBASE-CLIENTzfire-admin-python/{0})
credentialheadersbase_url)r   _project_idr   r	   PROJECT_URLr   Z_project_urlfirebase_admin__version__r   ZJsonHttpClientr   get_credential_clientOPERATION_URL_operation_client)r,   r   Z
ml_headersr   r   r   r-   B  s&     
z_MLService.__init__c              
   C   sP   t | z| jjd|dW S  tjjk
rJ } zt|W 5 d }~X Y nX d S )NrD   url)r   r  bodyrequestsr   RequestExceptionr   #handle_platform_error_from_requests)r,   rd   errorr   r   r   rc   U  s
    z_MLService.get_operationc                 C   s^   t tj|}|tj }|dk	rP|tj   }|dk rBtdt	||d }t
| dS )zGSleeps for the appropriate amount of time. Or throws deadline exceeded.Nr   zPolling max time exceeded.)powr	   POLL_EXPONENTIAL_BACKOFF_FACTORPOLL_BASE_WAIT_TIME_SECONDSrK   nowrO   r   ZDeadlineExceededErrormintimesleep)r,   current_attempt	stop_timeZdelay_factorZwait_time_secondsZmax_seconds_leftr   r   r   _exponential_backoff\  s    

z_MLService._exponential_backoffFNc                 C   sB  t |tstd|drV|dr0|dS |drJt|dtjddt|d}|di }|d	d
}|	dstdt
|d\}}d}	tj }
|dkrdn|
tj|d }|r|ds| |	| | |}|	d7 }	q|dr6|dr|dS |dr6t|dt| S )a  Handles long running operations.

        Args:
            operation: The operation to handle.
            wait_for_operation: Should we allow polling for the operation to complete.
                If no polling is requested, a locked model will be returned instead.
            max_time_seconds: The maximum seconds to try polling for operation complete.
                (None for no limit)

        Returns:
            dict: A dictionary of the returned model properties.

        Raises:
            TypeError: if the operation is not a dictionary.
            ValueError: If the operation is malformed.
            UnknownError: If the server responds with an unexpected response.
            err: If the operation exceeds polling attempts or stop_time
        zOperation must be a dictionary.doneresponser
  z$Internal Error: Malformed Operation.)rV   rC   metadatar0   r   ZModelOperationMetadataz#Unknown type of operation metadata.r   N)secondsr   )r<   r2   r   rD   r   Zhandle_operation_errorr   ZUnknownErrorr   endswithrE   rK   r  r   r  rc   r   rk   )r,   Z	operationr`   ra   rd   r  Zmetadata_typerF   r   r  
start_timer  r   r   r   rb   h  s:    









z_MLService.handle_operationc              
   C   s`   t | z"| | jjdd|jdddW S  tjjk
rZ } zt	|W 5 d }~X Y nX d S )Npostr   Trh   r  json)
r   rb   r  r  rk   r  r   r  r   r	  )r,   r   r
  r   r   r   r     s    z_MLService.create_modelc              
   C   s   t || d|j}|d k	r,|d| }z"| | jjd||jdddW S  tjj	k
r~ } zt
|W 5 d }~X Y nX d S )N
models/{0}z?updateMask={0}patchTrh   r  )r   r   r   rb   r  r  rk   r  r   r  r   r	  )r,   r   r   r   r
  r   r   r   r     s    
z_MLService.update_modelc                 C   s8   t | d| j|}t|d|id}| j|ddS )Nzprojects/{0}/models/{1}rX   )rC   rU   zstate.published)r   )r   r   r   r   r   r   )r,   r   r   Z
model_namer   r   r   r   r     s     z_MLService.set_publishedc              
   C   sV   t | z| jjdd|dW S  tjjk
rP } zt|W 5 d }~X Y nX d S )NrD   r  r  	r   r  r  r   r  r   r  r   r	  r,   r   r
  r   r   r   r     s
    z_MLService.get_modelc              
   C   s   t | t| t| i }|r(||d< |r4||d< |r@||d< d}|rhtt| d}|d | }z| jjd|dW S  t	j
jk
r } zt|W 5 d	}~X Y nX d	S )
z lists Firebase ML models.filterr"   r#   r   T?rD   r  N)r   r   r   r   	urlencodesorteditemsr  r  r  r   r  r   r	  )r,   r!   r"   r#   paramsr   Z	param_strr
  r   r   r   r      s$    z_MLService.list_modelsc              
   C   sX   t | z| jjdd|d W n0 tjjk
rR } zt|W 5 d }~X Y nX d S )Ndeleter  r  r   r!  r   r   r   r$     s
    z_MLService.delete_model)FN)N)rm   rn   ro   rp   r   r  r  r  r-   rc   r  rb   r   r   r   r   r    r$   r   r   r   r   r	   :  s   
9
r	   )N)N)N)N)N)NNNN)N)N);rp   rK   rer  r   urllibr   r   r  r   r   r   r   r   r   r   Z
tensorflowr   r   r   r   compiler   r   r   r   r   r   r   r   r   r   r   r   r   r    r$   r   ru   r4   r   r   r   r   r   r   rE   r   r   r   rJ   r\   r   r   rf   r   r   r   r	   r   r   r   r   <module>   s   











 &D/ #B

	
