
    	]jI                     T   d dl 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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 d dlmZ d d	lmZ d d
lmZ d dlmZ d dl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* dZ+dZ, ejZ                  e.      Z/ G d d      Z0d Z1d Z2y)    N)datetime)reduceflag_set)redis_connectedstart_job_async_or_sync)batch)SerializableGeneratorget_all_dirs_from_dirget_all_files_from_dirget_temp_dir)View)settings)File)temp)transaction)Prefetch)Q)
dateformattimezone)	Converter)
AnnotationAnnotationDraftTaskonlyexcludec                   z    e Zd Zd Zd ZddZddZed        Zd Z	ddZ
d	 Zed
        Zd ZddZddZddZy)ExportMixinc                 Z    | j                   |_         | j                   j                  |      S N)projecthas_permission)selfusers     I/root/env/lib/python3.12/site-packages/label_studio/data_export/mixins.pyr"   zExportMixin.has_permission%   s"    ||||**400    c                     | j                   j                  j                  dd       dt        j                  t        j                         d       S )N --at-z	Y-m-d-H-i)r!   titlereplacer   formatr   now)r#   s    r%   get_default_titlezExportMixin.get_default_title)   s@    ,,$$,,S#67tJ<M<Mhlln^i<j;kllr&   Nc                 J   t        |t              s|S d|v r	 t        |d         }t        j                  j                  | j                  |      j                  d      }t        j                  j                  |      j                  dd      }|j                  |      }d|v r<|d   }|t"        k(  r|j                  d      }n|t$        k(  r|j'                  d      }d|v r<|d   }|t"        k(  r|j                  d      }n|t$        k(  r|j'                  d      }d|v r=|d   }|t"        k(  r|j                  d      }|S |t$        k(  r|j'                  d      }|S # t        t        j                  f$ r"}t        j!                  d	|        Y d
}~d
}~ww xY w)a'  
        task_filter_options: None or Dict({
            view: optional int id or View
            skipped: optional None or str:("include|exclude")
            finished: optional None or str:("include|exclude")
            annotated: optional None or str:("include|exclude")
        })
        view)r!   idT)add_selected_items)prepare_paramsr2   flatid__inzIncorrect view params Nskipped)annotations__was_cancelledfinished)
is_labeled	annotatedF)
isinstancedictintr   objectsgetr!   get_prepare_tasks_paramsr   preparedonly_filteredvalues_listfilter
ValueErrorDoesNotExistloggerwarningONLYEXCLUDEr   )r#   taskstask_filter_optionsvaluer4   	tab_tasksexcs          r%   _get_filtered_taskszExportMixin._get_filtered_tasks,   s    -t4L((?/78!%!1!1$,,5!1!Q!j!j'+ "k " !MM77~7Vbbcgnrbs	I6 ++'	2E}E'!F,,'
3E}5'!6--'4E}F  '!G-  1 12 ?!7u=>>?s   BE' 'F" FF"c                 t   t         j                  j                         }t        |t              rg }|j                  d      r|j                  t        dd             |j                  d      r|j                  t        d             |j                  d      r|j                  t        d             |rt        d	 |      }|j                  |      }|j                  d
      }t        |j                  d      r;ddlm} |j                  t!        d|j                  j                  d                  }|S )a*  
        Filtering using disjunction of conditions

        annotation_filter_options: None or Dict({
            usual: optional None or bool:("true|false")
            ground_truth: optional None or bool:("true|false")
            skipped: optional None or bool:("true|false")
        })
        usualF)was_cancelledground_truthrW   T)rW   r9   )rV   c                     | |z  S r     )xys     r%   <lambda>z@ExportMixin._get_filtered_annotations_queryset.<locals>.<lambda>k   s
    A r&   completed_byreviewsr   )AnnotationReview
created_byqueryset)r   rA   allr>   r?   rB   appendr   r   rG   select_relatedhasattrmodelreviews.modelsr_   prefetch_relatedr   )r#   annotation_filter_optionsrb   q_listqr_   s         r%   "_get_filtered_annotations_querysetz.ExportMixin._get_filtered_annotations_querysetW   s     %%))+/6F(,,W5ae%HI(,,^<aT23(,,Y7ad34-v6#??1- **>:8>>9-700-=-E-E-T-TUa-bcH r&   c                 ,   dg i}t        | t              rd| v r;t        | d   t              r(| d   j                  d      s|d   j                  d       d| v r;t        | d   t              r(| d   j                  d      s|d   j                  d       d| v r(| d   j                  d      s|d   j                  d       dt        j
                  i|d<   d| v r| d   |d   d<   | j                  d	      d
u rdg|d<   | j                  d      du rd|d<   |S )Nexpanddraftsonly_idpredictionsannotations__completed_byzannotations.completed_byinterpolate_key_framescontextinclude_annotation_historyFzannotations.historyomitdownload_resourcesT)r>   r?   rB   rd   r   INTERPOLATE_KEY_FRAMES)serialization_optionsoptionss     r%   _get_export_serializer_optionz)ExportMixin._get_export_serializer_optionz   sF   R.+T2114X>E-h7;;IF!((2!664]CTJ-m<@@K!((7*.CCLa+Mc)nM !(()CD":H<[<[!\GI'+@@?TUm?n	"#;<$(()EF%O#8"9$(()=>$F04,-r&   c           	         ddl m} | j                  |      }t        | dd       } |d|      r& |d|      rt	        |d      r|j                         }t        j                  j                  |	      j                  d
      j                  t        d|      t        dt        j                  j                  d            d      } |d|      r& |d|      rt	        |d      r|j                         }|S )Nr   r   )rj   r`   *fflag_feat_fit_568_finite_state_management)r$   #fflag_feat_fit_710_fsm_state_fields
with_stater7   file_uploadannotationsra   rp   r$   comment_authors)core.feature_flagsr   rm   getattrrf   r   r   rA   rG   re   ri   r   r   )r#   idsrj   r   annotations_qsr$   qss          r%   get_task_querysetzExportMixin.get_task_queryset   s    /@@[t@u t\40AM>TJ5+668N LLs+^M*@O,C,C,R,RSY,Z[! 	 AM>TJL)B	r&   c              #     K   ddl m} t        j                  d       t	        j
                         }t        j                         5  ddi| _        | j                  j                  }t        j                  d       t        | j                  ||      j                         j                  dd	
            }| j                  |      }d}	t!        d| j                  j"                  j$                        r| j                  j'                         }
nt(        j*                  }
t-        ||
      D ]:  }|	dz  }	t        | j/                  ||            }t        j                  d|	|
z          t1        |t2              r;|j5                  d      r*|D cg c]  }|j6                  j9                         s|! }}|ro|j5                  d      d	u r\|D cg c]  }|j:                   }}t<        j>                  jA                  |      j                  dd	
      }| jC                  ||      } ||fdd	i|}| j                  dxx   tE        |      z  cc<   |jF                  D ]  }|  = 	 ddd       t	        j
                         |z
  }t        jI                  | j                  d    d| jJ                   d|jM                         dd       yc c}w c c}w # 1 sw Y   oxY ww)a+  
        serialization_options: None or Dict({
            drafts: optional
                None
                    or
                Dict({
                    only_id: true/false
                })
            predictions: optional
                None
                    or
                Dict({
                    only_id: true/false
                })
            annotations__completed_by: optional
                None
                    or
                Dict({
                    only_id: true/false
                })
        })
           )ExportDataSerializerzRun get_task_querysettask_numberr   zTasks filtration)rO   r2   Tr5   0fflag_fix_back_plt_807_batch_size_26062025_shortzBatch: only_with_annotationsrv   )task_id__inmanyNz tasks from project z exported in z.2fz seconds)'serializersr   rJ   debugr   r.   r   atomiccountersr!   rN   listrS   distinctrF   r|   r   organizationr`   get_task_batch_sizer   
BATCH_SIZEr	   r   r>   r?   rB   r   existsr2   r   rA   rG   update_export_serializer_optionlendatainfo
project_idtotal_seconds)r#   rO   rj   rz   r   start	all_taskstask_idsbase_export_serializer_optionir   r   rN   taskannotation_ids
serializerdurations                    r%   get_export_datazExportMixin.get_export_data   s    . 	6,-! $	 +A.DM**ILL+,((H[(\T-H
 -1,N,NOd,e)AJDLLLeLeLpLpq!\\==?
%00
Xz2 QT33C9RSTwq|n56148=P=T=TUl=m.3Qdt7G7G7N7N7PTQEQ(-B-F-FGc-dhl-l49:D:H:%/%7%7%>%>8%>%T%`%`aelp%`%qN484X4X5~51 2%dddFcd
m,E
:,&OO DJ!'$	J <<>E)}}]+,,@@QQ^_g_u_u_wx{^|  }E  F	
 R  ;7$	 $	sE   AK#EKK.K2KK BK)A$K#
KK K#c                     |S r    rY   )r#   r   r   s      r%   r   z+ExportMixin.update_export_serializer_option  s    ,,r&   c                     t        j                         }d|j                  z  }| j                  |      }|r%|j	                  |       | j                  |      }|r%|j                         }|S )N   )hashlibmd5
block_sizereadupdate	hexdigest)file
md5_objectr   chunkr   s        r%   eval_md5zExportMixin.eval_md5  sd    [[]
:000
		*%e$IIj)E  ""$
r&   c                 N   t        j                         }d| j                  j                   d|j	                  d       d|dd  d}| j                  j                   d| }t        ||	      }| j                  j                  ||       || _        | j                  g d
       y )Nzproject-r*   z%Y-%m-%d-%H-%Mr)   r      z.json/name)r   r   r   update_fields)	r   r.   r!   r2   strftimer   r   saver   )r#   r   r   r.   	file_name	file_pathfile_s          r%   	save_filezExportMixin.save_file  s    llnt||/tCLLAQ4R3SSTUXYZ[\U]T^^cd	||'q4	T	*		y%(		 ;	<r&   c                    t         j                  d| j                   d| d| d| d	       	 t        j                  d      j                  t        | j                  |||                  }t        j                  d	t        j                  
      5 }|D ]$  }|j                  d      }|j                  |       & |j                  d       | j                  |      }| j!                  ||       d d d        | j"                  j$                  | _        | j)                  dg       t1        j2                         | _        | j)                  dg       y # 1 sw Y   dxY w# t*        $ rN}	| j"                  j,                  | _        | j)                  dg       t         j/                  d|	       Y d }	~	d }	~	ww xY w# t1        j2                         | _        | j)                  dg       w xY w)NzRun export for z# with params:
task_filter_options: z
annotation_filter_options: z
serialization_options: 
F)ensure_asciirO   rj   rz   z.export.jsonsuffixdirzutf-8r   statusr   zExport was failed: %sfinished_at)rJ   r   r2   jsonJSONEncoder
iterencoder
   r   tempfileNamedTemporaryFiler   FILE_UPLOAD_TEMP_DIRencodewriteseekr   r   Status	COMPLETEDr   r   	ExceptionFAILED	exceptionr   r.   r   )
r#   rO   rj   rz   	iter_jsonr   r   encoded_chunkr   es
             r%   export_to_filezExportMixin.export_to_file  s   dggY '$$7#8 9**C)D E&&;%<B@	
	5((e<GG%((,?2K.C ) I ,,NHeHef *jn& .E$)LL$9MJJ}-. 		!mmD)tS)* ++//DKIIXJI/  (||~DII]OI4%* *  	9++,,DKIIXJI/4a88	9
  (||~DII]OI4sD   A%E  AE16E  EE   	F7)AF2-F: 2F77F: :.G(c           
      Z   | j                   | j                  j                  k(  rt        j	                  d       y | j                  j                  | _         | j                  dg       t               r&t        t        | j                  |||dt        d       y | j                  |||       y )Nz$Try to export with in progress stager   r   default3h)
queue_name
on_failurejob_timeoutr   )r   r   IN_PROGRESSrJ   rK   r   r   r   export_backgroundr2   set_export_background_failurer   )r#   rO   rj   rz   s       r%   run_file_exportingzExportMixin.run_file_exporting=  s    ;;$++111NNABkk--		
	+#!#)%$8 	 $7*C&;   r&   c           
      h	   t         j                  d| j                  | j                  |||       t	               5 }d}t        j                  |      |z  }|j                  ddd       t        | j                  j                         d ||| j                  j                  j                  j                  j                  |      }t        j                  | j                  j                         j                   }t        j                  |      |z  }	t         j                  d| j                  ||	       | j                  j#                         5 }
t#        |	d      5 }t%        j&                  |
|       d d d        d d d        t         j                  d	| j                  |	j)                         j*                         t         j                  d
| j                  ||       |j-                  |	||d       t         j                  d| j                  |       t/        |      }t1        |      }t         j                  d| j                  t3        |      t3        |             t3        |      dk(  r9t3        |      dk(  r+t         j                  d| j                  |       	 d d d        y t3        |      dk(  rUt3        |      dk(  rG|d   }t        j                  |      j4                  t        j                  |      j6                  z   }d}nit%        j8                  |d|       t        j                  |      t;        |j4                        dz   z  }t        j                  |      j4                  dz   }d}t         j                  d| j                  ||t        j                  |      j)                         j*                  |       t=        j>                  t        j                  |      j6                  t@        jB                        }t#        |d      5 }t%        j&                  ||       d d d        |jE                         }|jG                  d       t         j                  d| j                  ||       tI        ||      cd d d        S # 1 sw Y   xY w# 1 sw Y   xY w# 1 sw Y   |xY w# 1 sw Y   y xY w)NzeStarting export conversion: export_id=%s project_id=%s to_format=%s download_resources=%s hostname=%souti  T)modeparentsexist_ok)configproject_dir
upload_dirrx   access_tokenhostnamezMStaging export input for conversion: export_id=%s input_name=%s input_path=%swbz=Input staged for conversion: export_id=%s input_size_bytes=%sz:Running converter: export_id=%s to_format=%s output_dir=%sF)is_dirz-Converter finished: export_id=%s to_format=%sz;Conversion output discovered: export_id=%s files=%s dirs=%sr   z3Conversion has no output: export_id=%s to_format=%sr   single_filezipz.zipzip_archivez|Streaming conversion output to temporary file: export_id=%s output_kind=%s output_file=%s output_size_bytes=%s final_name=%sr   rb)r   z=Conversion file ready: export_id=%s filename=%s size_bytes=%sr   )%rJ   r   r2   r   r   pathlibPathmkdirr   r!   get_parsed_configr   r`   
auth_tokenkeyr   r   openshutilcopyfileobjstatst_sizeconvertr   r   r   stemr   make_archivestrr   r   r   r   tellr   r   )r#   	to_formatrx   r   tmp_dirOUTout_dir	converter
input_nameinput_file_pathfile_infile_outfilesdirsoutput_filefilenameoutput_kindresult_filefresult_sizes                       r%   convert_filezExportMixin.convert_fileW  s   A GGOO
	
 ^ W	4wCll7+c1GMMudTMB!||557 "#5!\\66AALLPP!I !diinn5::J%ll73j@OKK_	 ! 6Wd?D.I 6X""7H56 6KKO$$&.. KKL	 ow	%PKKGR[\*73E(1DKKME
D		 5zQ3t9>QSWSZSZ\efgW	4 W	4h UqSY!^#Ah"<<
3887<<;T;[;[[+##GUG<%ll73s7<<7H67QR"<<
3886A+KKH [)..088
 #55||H-4411K k- 3""1k23%**,KQKKO	 (3oW	4 W	4.6 6 6 6l3 3[W	4 W	4s]   DR(
RR.R6D!R( E*R(
R!AR(RRR	R(R%	!R((R1r    )NNN)FN)__name__
__module____qualname__r"   r/   rS   rm   staticmethodr|   r   r   r   r   r   r   r   r  rY   r&   r%   r   r   $   si    1m)V!F  :!FD
L-  =#5J4c4r&   r   c                 j    ddl m} |j                  j                  |       j	                  |||       y )Nr   Exportr2   )data_export.modelsr#  rA   rB   r   )	export_idrO   rj   rz   argskwargsr#  s          r%   r   r     s0     *
NN)$33!r&   c                     ddl m} | j                  d   }|j                  j	                  |      j                  |j                  j                         y )Nr   r"  r$  )r   )r%  r#  r'  rA   rG   r   r   r   )job
connectiontyperP   	tracebackr#  r&  s          r%   r   r     s?    )I
NNY'..fmm6J6J.Kr&   )3r   r   loggingr   r  r   	functoolsr   r   r   
core.redisr   r   core.utils.commonr	   core.utils.ior
   r   r   r   data_manager.modelsr   django.confr   django.core.filesr   r   r   	django.dbr   django.db.modelsr   django.db.models.query_utilsr   django.utilsr   r   label_studio_sdk.converterr   tasks.modelsr   r   r   rL   rM   	getLoggerr  rJ   r   r   r   rY   r&   r%   <module>r=     s           ' ? #  %   " . ! % * - 0 : :
 
		8	$V4 V4r	Lr&   