BiliOB/biliob_analyzer/add_keyword.py at 923a65b8591c9817189bcdbdb0df49d7db517086 · Jannchie/BiliOB · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
from pymongo import ReturnDocument
import jieba
from db import db
from time import sleep
# 载入字典
from biliob_tracer.task import ProgressTask


class KeywordAdder():

    def __init__(self):
        self.mongo_author = db['author']
        self.mongo_video = db['video']
        self.mongo_word = db['search_word']
        jieba.load_userdict('./biliob_analyzer/dict.txt')

    def get_video_kw_list(self, aid):
        # 关键字从name和official中提取
        video = self.mongo_video.find_one(
            {'aid': aid}, {'_id': 0, 'title': 1, 'channel': 1, 'subChannel': 1, 'author': 1, 'tag': 1})
        kw = []
        for each_key in video:
            if each_key != 'keyword' or each_key != 'tag':
                kw.append(str(video[each_key]).lower())
            elif each_key == 'tag':
                kw += video['tag']
            else:
                kw += video['keyword']
        seg_list = jieba.lcut_for_search(
            ' '.join(kw), True)  # 搜索引擎模式

        # 全名算作关键字
        if 'author' in video and video['author'].lower() not in seg_list:
            seg_list.append(video['author'].lower())

        while ' ' in seg_list:
            seg_list.remove(' ')
        while '、' in seg_list:
            seg_list.remove('、')
        return list(set(seg_list))

    def add_to_video(self, aid, seg_list):
        sleep(0.01)
        self.mongo_video.update_one({'aid': aid}, {'$set': {
            'keyword': seg_list
        }})

    def add_video_kw(self, aid):
        self.add_to_video(aid, self.get_video_kw_list(aid))
        return True

    def get_author_kw_list(self, mid):
        # 关键字从name和official中提取
        author = self.mongo_author.find_one(
            {'mid': mid}, {'_id': 0, 'name': 1, 'official': 1, 'keyword': 1})
        kw = []
        for each_key in author:
            if each_key != 'keyword':
                kw.append(str(author[each_key]).lower())
            else:
                kw += author['keyword']
        seg_list = jieba.lcut_for_search(
            ' '.join(kw), True)  # 搜索引擎模式

        # 全名算作关键字
        if 'name' in author and author['name'].lower() not in seg_list:
            seg_list.append(author['name'].lower())

        while ' ' in seg_list:
            seg_list.remove(' ')
        while '、' in seg_list:
            seg_list.remove('、')
        return list(set(seg_list))

    def add_author_kw(self, mid):
        self.add_to_author(mid, self.get_author_kw_list(mid))
        return True

    def add_to_author(self, mid, seg_list):
        sleep(0.01)
        self.mongo_author.update_one(
            {'mid': mid}, {'$set': {'keyword': seg_list}})

    def add_all_author(self):
        authors = self.mongo_author.find(
            {
                '$or': [
                    {
                        'keyword': []
                    }, {
                        'keyword': {
                            '$exists': False
                        }
                    }
                ]
            }, {'_id': 0, 'mid': 1}).batch_size(200)
        for each_author in authors:
            mid = each_author['mid']
            self.add_author_kw(mid)

    def add_all_video(self):
        videos = self.mongo_video.find({
            '$or': [
                {
                    'keyword': []
                }, {
                    'keyword': {
                        '$exists': False
                    }
                }
            ]
        }, {'_id': 0, 'aid': 1}).batch_size(200)
        for each_video in videos:
            aid = each_video['aid']
            self.add_video_kw(aid)

    def refresh_all_author(self):
        authors = self.mongo_author.find(
            {}, {'_id': 0, 'mid': 1}).batch_size(500)
        for each_author in authors:
            mid = each_author['mid']
            print("[mid]"+str(mid))
            self.add_author_kw(mid)

    def refresh_all_video(self):
        videos = self.mongo_video.find(
            {}, {'_id': 0, 'aid': 1}).batch_size(500)
        for each_video in videos:
            aid = each_video['aid']
            print("[aid]"+str(aid))
            self.add_video_kw(aid)

    def add_omitted(self):
        total_value = self.mongo_word.count_documents({})
        if self.mongo_word.count_documents({}) < 100:
            return
        t = ProgressTask("更新查询关键词字典", total_value=total_value,collection=db['tracer'])
        d = open('./biliob_analyzer/dict.txt', 'r',
                 encoding='utf8').read().split('\n')
        for each in self.mongo_word.find():
            if 'aid' in each and each['aid'] not in d:
                d.append(each['aid'])
            elif 'mid' in each and each['mid'] not in d:
                d.append(each['mid'])
            t.current_value += 1
        pass
        t.finished = True
        o = open('./biliob_analyzer/dict.txt',
                 'w', encoding='utf8', newline='')
        for each in d:
            o.write(each+'\n')
        o.close()
        self.mongo_word.delete_many({})
        jieba.load_userdict('./biliob_analyzer/dict.txt')
        self.refresh_all_author()
        self.refresh_all_video()