← Back to branch summary

~ubuntu-branches/ubuntu/karmic/spambayes/karmic

« back to all changes in this revision

Viewing changes to spambayes/msgs.py

  • Committer: Bazaar Package Importer
  • Author(s): Jorge Bernal
  • Date: 2005-04-07 14:02:02 UTC
  • Revision ID: james.westby@ubuntu.com-20050407140202-mgyh6t7gn2dlrrw5
Tags: upstream-1.0.1
Import upstream version 1.0.1

Show diffs side-by-side

added added

removed removed

Lines of Context:
spambayes/msgs.py
 
1
from __future__ import generators
 
2
 
 
3
import os
 
4
import random
 
5
 
 
6
from spambayes.tokenizer import tokenize
 
7
 
 
8
HAMTEST  = None
 
9
SPAMTEST = None
 
10
HAMTRAIN  = None
 
11
SPAMTRAIN = None
 
12
SEED = random.randrange(2000000000)
 
13
 
 
14
class Msg(object):
 
15
    __slots__ = 'tag', 'guts'
 
16
 
 
17
    def __init__(self, dir, name):
 
18
        path = dir + "/" + name
 
19
        self.tag = path
 
20
        f = open(path, 'rb')
 
21
        self.guts = f.read()
 
22
        f.close()
 
23
 
 
24
    def __iter__(self):
 
25
        return tokenize(self.guts)
 
26
 
 
27
    # Compare msgs by their paths; this is appropriate for sets of msgs.
 
28
    def __hash__(self):
 
29
        return hash(self.tag)
 
30
 
 
31
    def __eq__(self, other):
 
32
        return self.tag == other.tag
 
33
 
 
34
    def __str__(self):
 
35
        return self.guts
 
36
 
 
37
# The iterator yields a stream of Msg objects, taken from a list of
 
38
# directories.
 
39
class MsgStream(object):
 
40
    __slots__ = 'tag', 'directories', 'keep'
 
41
 
 
42
    def __init__(self, tag, directories, keep=None):
 
43
        self.tag = tag
 
44
        self.directories = directories
 
45
        self.keep = keep
 
46
 
 
47
    def __str__(self):
 
48
        return self.tag
 
49
 
 
50
    def produce(self):
 
51
        if self.keep is None:
 
52
            for directory in self.directories:
 
53
                for fname in os.listdir(directory):
 
54
                    yield Msg(directory, fname)
 
55
            return
 
56
        # We only want part of the msgs.  Shuffle each directory list, but
 
57
        # in such a way that we'll get the same result each time this is
 
58
        # called on the same directory list.
 
59
        for directory in self.directories:
 
60
            all = os.listdir(directory)
 
61
            random.seed(hash(max(all)) ^ SEED) # reproducible across calls
 
62
            random.shuffle(all)
 
63
            del all[self.keep:]
 
64
            all.sort()  # seems to speed access on Win98!
 
65
            for fname in all:
 
66
                yield Msg(directory, fname)
 
67
 
 
68
    def __iter__(self):
 
69
        return self.produce()
 
70
 
 
71
class HamStream(MsgStream):
 
72
    def __init__(self, tag, directories, train=0):
 
73
        if train:
 
74
            MsgStream.__init__(self, tag, directories, HAMTRAIN)
 
75
        else:
 
76
            MsgStream.__init__(self, tag, directories, HAMTEST)
 
77
 
 
78
class SpamStream(MsgStream):
 
79
    def __init__(self, tag, directories, train=0):
 
80
        if train:
 
81
            MsgStream.__init__(self, tag, directories, SPAMTRAIN)
 
82
        else:
 
83
            MsgStream.__init__(self, tag, directories, SPAMTEST)
 
84
 
 
85
def setparms(hamtrain, spamtrain, hamtest=None, spamtest=None, seed=None):
 
86
    """Set HAMTEST/TRAIN and SPAMTEST/TRAIN.
 
87
       If seed is not None, also set SEED.
 
88
       If (ham|spam)test are not set, set to the same as the (ham|spam)train
 
89
       numbers (backwards compat option).
 
90
    """
 
91
 
 
92
    global HAMTEST, SPAMTEST, HAMTRAIN, SPAMTRAIN, SEED
 
93
    HAMTRAIN, SPAMTRAIN = hamtrain, spamtrain
 
94
    if hamtest is None:
 
95
        HAMTEST = HAMTRAIN
 
96
    else:
 
97
        HAMTEST = hamtest
 
98
    if spamtest is None:
 
99
        SPAMTEST = SPAMTRAIN
 
100
    else:
 
101
        SPAMTEST = spamtest
 
102
    if seed is not None:
 
103
        SEED = seed