Skip to content
Snippets Groups Projects
search.py 51 KiB
Newer Older
  • Learn to ignore specific revisions
  • Vilem Duha's avatar
    Vilem Duha committed
                    killthreads_sml.append(k)  # do actual killing here?
    
            killthreads_full = []
            for k in thumb_full_download_threads.keys():
                if k not in thumb_full_filepaths:
                    killthreads_full.append(k)  # do actual killing here?
            # TODO do the killing/ stopping here! remember threads might have finished inbetween!
    
            if self.stopped():
    
    Vilém Duha's avatar
    Vilém Duha committed
                utils.p('stopping search : ' + str(query))
    
    Vilem Duha's avatar
    Vilem Duha committed
                return
    
            # this loop handles downloading of small thumbnails
            for imgpath, url in sml_thbs:
                if imgpath not in thumb_sml_download_threads and not os.path.exists(imgpath):
                    thread = ThumbDownloader(url, imgpath)
                    # thread = threading.Thread(target=download_thumbnail, args=([url, imgpath]),
                    #                           daemon=True)
                    thread.start()
                    thumb_sml_download_threads[imgpath] = thread
                    # threads.append(thread)
    
                    if len(thumb_sml_download_threads) > maxthreads:
                        while len(thumb_sml_download_threads) > maxthreads:
                            threads_copy = thumb_sml_download_threads.copy()  # because for loop can erase some of the items.
                            for tk, thread in threads_copy.items():
                                if not thread.is_alive():
                                    thread.join()
    
    Vilem Duha's avatar
    Vilem Duha committed
                                    # utils.p(x)
    
    Vilem Duha's avatar
    Vilem Duha committed
                                    del (thumb_sml_download_threads[tk])
    
    Vilem Duha's avatar
    Vilem Duha committed
                                    # utils.p('fetched thumbnail ', i)
    
    Vilem Duha's avatar
    Vilem Duha committed
                                    i += 1
            if self.stopped():
    
    Vilém Duha's avatar
    Vilém Duha committed
                utils.p('stopping search : ' + str(query))
    
    Vilem Duha's avatar
    Vilem Duha committed
                return
            idx = 0
            while len(thumb_sml_download_threads) > 0:
                threads_copy = thumb_sml_download_threads.copy()  # because for loop can erase some of the items.
                for tk, thread in threads_copy.items():
                    if not thread.is_alive():
                        thread.join()
                        del (thumb_sml_download_threads[tk])
                        i += 1
    
            if self.stopped():
    
    Vilém Duha's avatar
    Vilém Duha committed
                utils.p('stopping search : ' + str(query))
    
    Vilem Duha's avatar
    Vilem Duha committed
                return
    
            # start downloading full thumbs in the end
            for imgpath, url in full_thbs:
                if imgpath not in thumb_full_download_threads and not os.path.exists(imgpath):
                    thread = ThumbDownloader(url, imgpath)
                    # thread = threading.Thread(target=download_thumbnail, args=([url, imgpath]),
                    #                           daemon=True)
                    thread.start()
                    thumb_full_download_threads[imgpath] = thread
            mt('thumbnails finished')
    
    
    def build_query_common(query, props):
    
        '''add shared parameters to query'''
        query_common = {}
    
    Vilém Duha's avatar
    Vilém Duha committed
        if props.search_keywords != '':
    
            query_common["query"] = props.search_keywords
    
        if props.search_verification_status != 'ALL':
            query_common['verification_status'] = props.search_verification_status.lower()
    
        if props.search_file_size:
    
            query_common["files_size_gte"] = props.search_file_size_min * 1024 * 1024
            query_common["files_size_lte"] = props.search_file_size_max * 1024 * 1024
    
        query.update(query_common)
    
    Vilem Duha's avatar
    Vilem Duha committed
    
    
    def build_query_model():
        '''use all search input to request results from server'''
    
        props = bpy.context.scene.blenderkit_models
        query = {
            "asset_type": 'model',
    
            # "engine": props.search_engine,
            # "adult": props.search_adult,
    
    Vilem Duha's avatar
    Vilem Duha committed
        }
        if props.search_style != 'ANY':
            if props.search_style != 'OTHER':
    
                query["model_style"] = props.search_style
    
    Vilem Duha's avatar
    Vilem Duha committed
            else:
    
                query["model_style"] = props.search_style_other
    
        # if props.search_advanced:
        if props.search_condition != 'UNSPECIFIED':
            query["condition"] = props.search_condition
        if props.search_design_year:
            query["designYear_gte"] = props.search_design_year_min
            query["designYear_lte"] = props.search_design_year_max
        if props.search_polycount:
            query["faceCount_gte"] = props.search_polycount_min
            query["faceCount_lte"] = props.search_polycount_max
    
        if props.search_texture_resolution:
            query["textureResolutionMax_gte"] = props.search_texture_resolution_min
            query["textureResolutionMax_lte"] = props.search_texture_resolution_max
    
    Vilém Duha's avatar
    Vilém Duha committed
    
    
    Vilem Duha's avatar
    Vilem Duha committed
        build_query_common(query, props)
    
        return query
    
    
    def build_query_scene():
        '''use all search input to request results from server'''
    
        props = bpy.context.scene.blenderkit_scene
        query = {
            "asset_type": 'scene',
    
            # "engine": props.search_engine,
    
    Vilem Duha's avatar
    Vilem Duha committed
            # "adult": props.search_adult,
        }
        build_query_common(query, props)
        return query
    
    
    def build_query_material():
        props = bpy.context.scene.blenderkit_mat
        query = {
            "asset_type": 'material',
    
        }
    
        # if props.search_engine == 'NONE':
        #     query["engine"] = ''
        # if props.search_engine != 'OTHER':
        #     query["engine"] = props.search_engine
        # else:
        #     query["engine"] = props.search_engine_other
    
    Vilem Duha's avatar
    Vilem Duha committed
        if props.search_style != 'ANY':
            if props.search_style != 'OTHER':
                query["style"] = props.search_style
            else:
                query["style"] = props.search_style_other
    
        if props.search_procedural == 'TEXTURE_BASED':
            # todo this procedural hack should be replaced with the parameter
            query["textureResolutionMax_gte"] = 0
            # query["procedural"] = False
            if props.search_texture_resolution:
                query["textureResolutionMax_gte"] = props.search_texture_resolution_min
                query["textureResolutionMax_lte"] = props.search_texture_resolution_max
    
    
    
        elif props.search_procedural == "PROCEDURAL":
            # todo this procedural hack should be replaced with the parameter
            query["files_size_lte"] = 1024 * 1024
            # query["procedural"] = True
    
    Vilém Duha's avatar
    Vilém Duha committed
    
    
    Vilem Duha's avatar
    Vilem Duha committed
        build_query_common(query, props)
    
        return query
    
    
    def build_query_texture():
        props = bpy.context.scene.blenderkit_tex
        query = {
            "asset_type": 'texture',
    
        }
    
        if props.search_style != 'ANY':
            if props.search_style != 'OTHER':
                query["search_style"] = props.search_style
            else:
                query["search_style"] = props.search_style_other
    
        build_query_common(query, props)
    
        return query
    
    
    def build_query_brush():
        props = bpy.context.scene.blenderkit_brush
    
        brush_type = ''
        if bpy.context.sculpt_object is not None:
            brush_type = 'sculpt'
    
        elif bpy.context.image_paint_object:  # could be just else, but for future p
            brush_type = 'texture_paint'
    
        query = {
            "asset_type": 'brush',
    
    
            "mode": brush_type
    
    Vilem Duha's avatar
    Vilem Duha committed
        }
    
        build_query_common(query, props)
    
        return query
    
    
    def mt(text):
        global search_start_time, prev_time
        alltime = time.time() - search_start_time
        since_last = time.time() - prev_time
        prev_time = time.time()
    
    Vilem Duha's avatar
    Vilem Duha committed
        utils.p(text, alltime, since_last)
    
    Vilém Duha's avatar
    Vilém Duha committed
    def add_search_process(query, params, orig_result):
    
    Vilem Duha's avatar
    Vilem Duha committed
        global search_threads
    
        while (len(search_threads) > 0):
            old_thread = search_threads.pop(0)
            old_thread[0].stop()
            # TODO CARE HERE FOR ALSO KILLING THE THREADS...AT LEAST NOW SEARCH DONE FIRST WON'T REWRITE AN OLDER ONE
    
        tempdir = paths.get_temp_dir('%s_search' % query['asset_type'])
    
    Vilém Duha's avatar
    Vilém Duha committed
        thread = Searcher(query, params, orig_result)
    
    Vilem Duha's avatar
    Vilem Duha committed
        thread.start()
    
    
    Vilém Duha's avatar
    Vilém Duha committed
        search_threads.append([thread, tempdir, query['asset_type'], {}])  # 4th field is for results
    
    Vilém Duha's avatar
    Vilém Duha committed
        mt('search thread started')
    
    
    def get_search_simple(parameters, filepath=None, page_size=100, max_results=100000000, api_key=''):
        '''
        Searches and returns the
    
    
        Parameters
        ----------
        parameters - dict of blenderkit elastic parameters
        filepath - a file to save the results. If None, results are returned
        page_size - page size for retrieved results
        max_results - max results of the search
        api_key - BlenderKit api key
    
        Returns
        -------
        Returns search results as a list, and optionally saves to filepath
    
        '''
        headers = utils.get_headers(api_key)
        url = paths.get_api_url() + 'search/'
        requeststring = url + '?query='
        for p in parameters.keys():
            requeststring += f'+{p}:{parameters[p]}'
    
        requeststring += '&page_size=' + str(page_size)
        response = rerequests.get(requeststring, headers=headers)  # , params = rparameters)
        # print(r.json())
        search_results = response.json()
    
        results = []
        results.extend(search_results['results'])
        page_index = 2
        page_count = math.ceil(search_results['count'] / page_size)
        while search_results.get('next') and len(results) < max_results:
            print(f'getting page {page_index} , total pages {page_count}')
            response = rerequests.get(search_results['next'], headers=headers)  # , params = rparameters)
            search_results = response.json()
            # print(search_results)
            results.extend(search_results['results'])
            page_index += 1
    
        if not filepath:
            return results
    
        with open(filepath, 'w') as s:
            json.dump(results, s)
        print(f'retrieved {len(results)} assets from elastic search')
        return results
    
    def search(category='', get_next=False, author_id=''):
    
    Vilem Duha's avatar
    Vilem Duha committed
        ''' initialize searching'''
        global search_start_time
    
        user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
    
    
    Vilem Duha's avatar
    Vilem Duha committed
        search_start_time = time.time()
    
    Vilem Duha's avatar
    Vilem Duha committed
        scene = bpy.context.scene
    
        ui_props = scene.blenderkitUI
    
        if ui_props.asset_type == 'MODEL':
    
    Vilem Duha's avatar
    Vilem Duha committed
            if not hasattr(scene, 'blenderkit'):
                return;
            props = scene.blenderkit_models
            query = build_query_model()
    
    
        if ui_props.asset_type == 'SCENE':
    
    Vilem Duha's avatar
    Vilem Duha committed
            if not hasattr(scene, 'blenderkit_scene'):
                return;
            props = scene.blenderkit_scene
            query = build_query_scene()
    
        if ui_props.asset_type == 'MATERIAL':
    
    Vilem Duha's avatar
    Vilem Duha committed
            if not hasattr(scene, 'blenderkit_mat'):
                return;
    
    Vilem Duha's avatar
    Vilem Duha committed
            props = scene.blenderkit_mat
            query = build_query_material()
    
        if ui_props.asset_type == 'TEXTURE':
    
    Vilem Duha's avatar
    Vilem Duha committed
            if not hasattr(scene, 'blenderkit_tex'):
                return;
            # props = scene.blenderkit_tex
            # query = build_query_texture()
    
    
        if ui_props.asset_type == 'BRUSH':
    
    Vilem Duha's avatar
    Vilem Duha committed
            if not hasattr(scene, 'blenderkit_brush'):
                return;
            props = scene.blenderkit_brush
            query = build_query_brush()
    
        if props.is_searching and get_next == True:
            return;
    
    Vilem Duha's avatar
    Vilem Duha committed
        if category != '':
    
            query['category_subtree'] = category
    
        if author_id != '':
    
            query['author_id'] = author_id
    
    
    Vilém Duha's avatar
    Vilém Duha committed
        elif props.own_only:
            # if user searches for [another] author, 'only my assets' is invalid. that's why in elif.
            profile = bpy.context.window_manager.get('bkit profile')
            if profile is not None:
                query['author_id'] = str(profile['user']['id'])
    
    
    Vilem Duha's avatar
    Vilem Duha committed
        # utils.p('searching')
    
    Vilem Duha's avatar
    Vilem Duha committed
        props.is_searching = True
    
        params = {
            'scene_uuid': bpy.context.scene.get('uuid', None),
            'addon_version': version_checker.get_addon_version(),
    
            'api_key': user_preferences.api_key,
    
    Vilem Duha's avatar
    Vilem Duha committed
            'get_next': get_next
        }
    
    
        # if free_only:
        #     query['keywords'] += '+is_free:true'
    
    Vilém Duha's avatar
    Vilém Duha committed
        orig_results = scene.get(f'bkit {ui_props.asset_type.lower()} search orig', {})
        if orig_results != {}:
    
    Vilém Duha's avatar
    Vilém Duha committed
            # ensure it's a copy in dict for what we are passing to thread:
    
    Vilém Duha's avatar
    Vilém Duha committed
            orig_results = orig_results.to_dict()
        add_search_process(query, params, orig_results)
    
        tasks_queue.add_task((ui.add_report, ('BlenderKit searching....', 2)))
    
    Vilem Duha's avatar
    Vilem Duha committed
        props.report = 'BlenderKit searching....'
    
    
    def search_update(self, context):
    
    Vilem Duha's avatar
    Vilem Duha committed
        utils.p('search updater')
    
        # if self.search_keywords != '':
    
        ui_props = bpy.context.scene.blenderkitUI
        if ui_props.down_up != 'SEARCH':
            ui_props.down_up = 'SEARCH'
    
    
    Vilém Duha's avatar
    Vilém Duha committed
        # here we tweak the input if it comes form the clipboard. we need to get rid of asset type and set it to
        sprops = utils.get_search_props()
        instr = 'asset_base_id:'
        atstr = 'asset_type:'
        kwds = sprops.search_keywords
        idi = kwds.find(instr)
        ati = kwds.find(atstr)
        # if the asset type already isn't there it means this update function
        # was triggered by it's last iteration and needs to cancel
    
    Vilém Duha's avatar
    Vilém Duha committed
        if idi > -1 and ati == -1:
    
    Vilém Duha's avatar
    Vilém Duha committed
            return;
        if ati > -1:
            at = kwds[ati:].lower()
            # uncertain length of the remaining string -  find as better method to check the presence of asset type
            if at.find('model') > -1:
                ui_props.asset_type = 'MODEL'
            elif at.find('material') > -1:
                ui_props.asset_type = 'MATERIAL'
            elif at.find('brush') > -1:
                ui_props.asset_type = 'BRUSH'
            # now we trim the input copypaste by anything extra that is there,
            # this is also a way for this function to recognize that it already has parsed the clipboard
            # the search props can have changed and this needs to transfer the data to the other field
            # this complex behaviour is here for the case where the user needs to paste manually into blender?
            sprops = utils.get_search_props()
            sprops.search_keywords = kwds[:ati].rstrip()
    
    Vilem Duha's avatar
    Vilem Duha committed
    
    
    class SearchOperator(Operator):
        """Tooltip"""
        bl_idname = "view3d.blenderkit_search"
        bl_label = "BlenderKit asset search"
    
        bl_description = "Search online for assets"
    
    Vilém Duha's avatar
    Vilém Duha committed
        bl_options = {'REGISTER', 'UNDO', 'INTERNAL'}
    
    Vilem Duha's avatar
    Vilem Duha committed
        own: BoolProperty(name="own assets only",
                          description="Find all own assets",
                          default=False)
    
    Vilem Duha's avatar
    Vilem Duha committed
        category: StringProperty(
            name="category",
            description="search only subtree of this category",
    
            options={'SKIP_SAVE'}
    
        author_id: StringProperty(
            name="Author ID",
            description="Author ID - search only assets by this author",
    
            options={'SKIP_SAVE'}
    
    Vilem Duha's avatar
    Vilem Duha committed
        get_next: BoolProperty(name="next page",
                               description="get next page from previous search",
    
                               options={'SKIP_SAVE'}
                               )
    
        keywords: StringProperty(
    
            name="Keywords",
            description="Keywords",
    
            options={'SKIP_SAVE'}
    
    Vilem Duha's avatar
    Vilem Duha committed
        @classmethod
        def poll(cls, context):
            return True
    
        def execute(self, context):
    
            # TODO ; this should all get transferred to properties of the search operator, so sprops don't have to be fetched here at all.
            sprops = utils.get_search_props()
            if self.author_id != '':
                sprops.search_keywords = ''
            if self.keywords != '':
                sprops.search_keywords = self.keywords
    
            search(category=self.category, get_next=self.get_next, author_id=self.author_id)
    
            # bpy.ops.view3d.blenderkit_asset_bar()
    
    Vilem Duha's avatar
    Vilem Duha committed
    
            return {'FINISHED'}
    
    
    classes = [
        SearchOperator
    ]
    
    
    def register_search():
        bpy.app.handlers.load_post.append(scene_load)
    
        for c in classes:
            bpy.utils.register_class(c)
    
    
        user_preferences = bpy.context.preferences.addons['blenderkit'].preferences
        if user_preferences.use_timers:
            bpy.app.timers.register(timer_update)
    
    Vilem Duha's avatar
    Vilem Duha committed
    
    
    def unregister_search():
        bpy.app.handlers.load_post.remove(scene_load)
    
        for c in classes:
            bpy.utils.unregister_class(c)
    
        if bpy.app.timers.is_registered(timer_update):
            bpy.app.timers.unregister(timer_update)