utils/spider.py at master · codeunsolved/utils · 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
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# PROGRAM : spider
# AUTHOR  : codeunsolved@gmail.com
# CREATED : February 15 2017
# VERSION : v0.0.1
# UPDAYE  : [v0.0.1] May 18 2018
# 1. optimize :AsyncSpider:'s INPUT and structure;

import time
from collections import defaultdict

from tornado import httpclient, gen, ioloop, queues

from .base import color_term


class AsyncSpider(object):
    """AsyncSpider.

    :param unfetched: dict{label: url}
    :param handle_response: func(response, label, url)

    :Refer:
    [example - a concurrent web spider](http://www.tornadoweb.org/en/stable/guide/queues.html?highlight=spider)
    """

    def __init__(self, unfetched,
                 handle_response,
                 concurrency=3, max_trial=10):
        self.unfetched = unfetched
        self.labels = sorted(self.unfetched.keys())
        self.handle_response = handle_response

        self.q = queues.Queue()
        self.fetching, self.fetched = set(), set()

        self.concurrency = concurrency
        self.max_trial = max_trial

        self.failed_trials = defaultdict(int)
        self.abandoned_trials = []

        self.start = None

        self.check_handle_response()
        [self.q.put(x) for x in self.labels]

    def check_handle_response(self):
        if not hasattr(self.handle_response, '__call__'):
            color_term("`handle_response()` is invalid!", 'ERR', exit(1))

    @gen.coroutine
    def get(self, label, url):
        try:
            resp = yield httpclient.AsyncHTTPClient().fetch(url)
        except Exception as e:
            color_term("Exception: {} for '{}'({})".format(e, label, url), 'ERR')
            raise gen.Return([])
        else:
            raise gen.Return(resp)

    @gen.coroutine
    def query(self):
        @gen.coroutine
        def fetch(i):
            label = yield self.q.get()
            url = self.unfetched[label]
            try:
                color_term("- [{}/{}]@{} {}".format(self.labels.index(label) + 1, len(self.labels), i, label))
                if label in self.fetching:
                    color_term("{} is fetching!".format(label), 'ERR')
                else:
                    self.fetching.add(label)
                    response = yield self.get(label, url)
                    if response:
                        self.handle_response(response, label, url)
                        self.fetched.add(label)
                    else:
                        self.failed_trials[label] += 1
                        if self.failed_trials[label] < self.max_trial:
                            color_term("FAILED! put {} back to the queue".format(label), 'WRN')
                            self.fetching.remove(label)
                            self.q.put(label)
                        else:
                            color_term("{} have reached max trails({})".format(label, self.max_trial), 'WRN')
                            self.fetched.add(label)
                            self.abandoned_trials.append(label)
            finally:
                self.q.task_done()

        @gen.coroutine
        def worker(i):
            while True:
                yield fetch(i)

        for _ in range(self.concurrency):
            worker(_)

        yield self.q.join()
        assert self.fetching == self.fetched
        print('------------------')

        if self.unfetched:
            color_term("Done in {:.0f} min {:.2f} sec, fetched {} URLs".format(
                *(divmod(time.time() - self.start, 60) + (len(self.unfetched),))), 'green')
        else:
            color_term("Nothing need to fetch")

        if self.abandoned_trials:
            print('------------------')
            color_term("Abandoned trials:", 'ERR')
            for x in self.abandoned_trials:
                print(x)

    def run(self):
        color_term("• start AsyncSpider")
        self.start = time.time()
        io_loop = ioloop.IOLoop.current()
        io_loop.run_sync(self.query)