<commit_before># -*- coding: utf-8  -*-
from dal import autocomplete
from django.core.exceptions import PermissionDenied
from django.core.urlresolvers import reverse
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.db.models import Count
from django.http import Http404, HttpResponseRedirect
from django.shortcuts import render, get_object_or_404
from django.views.decorators import csrf
from .forms import NominationForm
from .models import SACYear, Position, Nomination, \
                    NominationAnnouncement, VoteNomination, \
                    city_choices
from . import decorators, utils
import accounts.utils
import re

def show_index(request):
    if request.user.is_authenticated():
        return render(request, 'accounts/home.html')
    else:
        return render(request, 'accounts/home_unauthenticated.html')

@login_required
def list_positions(request, entity):
    current_year = SACYear.objects.get_current()
    if not current_year.is_nomination_open():
        return render(request, "voting/closed.html")

    user_nominations = Nomination.objects.filter(position__entity=entity,
                                                 position__year=current_year,
                                                 user=request.user)
    positions = Position.objects.filter(entity=entity, year=current_year)
    if not request.user.is_superuser:
        positions = positions.filter(colleges_allowed_to_nominate=request.user.profile.college)

    context = {'user_nominations': user_nominations,
               'positions': positions}
    return render(request, 'voting/list_' + entity + '_positions.html', context)

@login_required
def add_nominee(request, position_id):
    current_year = SACYear.objects.get_current()
    if not current_year.is_nomination_open():
        return HttpResponseRedirect(reverse("voting:closed"))

    position = get_object_or_404(Position, pk=position_id)
    context = {'position': position}

    user_nominations = Nomination.objects.filter(position__entity=position.entity,
                                                 position__year=current_year,
                                                 user=request.user)
    if user_nominations.exists():
        context['already_on'] = True
    else:
        if request.user.is_superuser:
            positions = Position.objects.filter(year=current_year)
        else:
            positions = Position.objects.filter(colleges_allowed_to_nominate=request.user.profile.college,
                                                year=current_year)
            if position not in positions:
                raise PermissionDenied
        if request.method == 'POST':
            instance = Nomination(user=request.user, position=position)
            form = NominationForm(request.POST, request.FILES, instance=instance)
            if form.is_valid():
                instance = form.save()
                return HttpResponseRedirect(reverse("voting:nomination_thanks", args=(position.pk,)))
        elif request.method == 'GET':
            form = NominationForm()
        context['form'] = form

    return render(request,'voting/add_nominee.html', context)

@login_required
def show_nomination(request, position_id, nomination_id):
    nomination = get_object_or_404(Nomination, position__pk=position_id,
                                   pk=nomination_id)
    if not request.user.is_superuser and \
       not nomination.user == request.user:
        raise PermissionDenied
    context = {'nomination': nomination}
    return render(request,'voting/show_nomination.html', context)

def announce_nominees(request, entity):
    current_year = SACYear.objects.get_current()
    context = {}
    if current_year.is_announcement_due():
        per_city = []
        for city_code, city_name in city_choices:
            positions = Position.objects\
                               .filter(entity=entity, city=city_code, year=current_year)\
                               .annotate(nomination_count=Count('nominationannouncement'))\
                               .filter(nomination_count__gt=1)
            if positions.exists():
                per_city.append((city_name, positions))
        context = {'entity': entity, 'per_city': per_city}

    return render(request, 'voting/announce_nominees.html', context)

@login_required
def show_voting_index(request):
    current_year = SACYear.objects.get_current()
    if not request.user.is_superuser and not current_year.is_voting_open():
        return render(request, "voting/voting_closed.html")
    else:
        position_pool = Position.objects.filter(year=current_year)\
                                        .annotate(announced_count=Count('nominationannouncement'))\
                                        .filter(announced_count__gte=2)\
                                        .exclude(votenomination__user=request.user)\
                                        .order_by('entity')
        if request.user.profile:
            position_pool = position_pool.filter(colleges_allowed_to_vote=request.user.profile.college)

        if position_pool.exists():
            completed_voting = False
            votes = VoteNomination.objects.none()
        else:
            completed_voting = True
            votes = VoteNomination.objects.filter(user=request.user)
        qrcode_value = utils.get_ticket(request.user)
        context = {'qrcode_value': qrcode_value,
                   'completed_voting': completed_voting,
                   'votes': votes}
        return render(request,'voting/show_voting.html', context)

@login_required
@decorators.ajax_only
@decorators.post_only
@csrf.csrf_exempt
def handle_vote(request):
    current_year = SACYear.objects.get_current()
    if not request.user.is_superuser and not current_year.is_voting_open():
        raise Exception("التصويت غير مُتاح حاليًا!")

    nomination_vote_pk = request.POST.get('nomination_vote_pk', None)
    if nomination_vote_pk:
        if nomination_vote_pk == 'skip':
            position_pk = request.POST.get('position_pk')
            nomination_announcement = None
            position = Position.objects.get(pk=position_pk)
        else:
            nomination_announcement = NominationAnnouncement.objects.get(pk=nomination_vote_pk)
            position = nomination_announcement.position

        previous_vote = VoteNomination.objects.filter(position=position,
                                                      user=request.user).exists()
        if previous_vote:
            raise Exception(u'سبق أن صوتّ لهذا المنصب')

        if not request.user.is_superuser and \
           not request.user.profile.college in\
               position.colleges_allowed_to_vote.all():
            raise PermissionDenied
        else:
            VoteNomination.objects.create(nomination_announcement=nomination_announcement,
                                          position=position,
                                          user=request.user)

    position_pool = Position.objects.filter(year=current_year)\
                                    .annotate(announced_count=Count('nominationannouncement'))\
                                    .filter(announced_count__gte=2)\
                                    .exclude(votenomination__user=request.user)
    if request.user.is_superuser:
        next_position = position_pool.first()
    else:
        next_position = position_pool.filter(colleges_allowed_to_vote=request.user.profile.college)\
                                     .first()
    if next_position:
        nominations = []
        for nomination in next_position.nominationannouncement_set.order_by('user__profile__ar_first_name'):
            nomination = {'pk': nomination.pk,
                          'nominee_name': nomination.user.profile.get_ar_full_name(),
                          'plan': nomination.plan.url,
                          'cv': nomination.cv.url}
            nominations.append(nomination)
        return {"position_title": next_position.title,
                "position_pk": next_position.pk,
                "entity": next_position.entity,
                "note": next_position.note,
                "nominations": nominations}
    else:
        return {'done': 1}

@login_required
def indicators(request):
    if not request.user.is_superuser:
        raise PermissionDenied
    current_year = SACYear.objects.get_current()
    election_positions = Position.objects.filter(year=current_year)\
                                         .annotate(announced_count=Count('nominationannouncement'))\
                                         .filter(announced_count__gte=2).order_by('entity', 'city')
    context = {'election_positions': election_positions}

    return render(request, 'voting/indicators.html', context)

@login_required
def list_votes_per_position(request, position_id):
    if not request.user.is_superuser:
        raise PermissionDenied

    position = get_object_or_404(Position, pk=position_id)
    votes = VoteNomination.objects.filter(position=position)
    context = {'position': position, 'votes': votes}

    return render(request, 'voting/list_votes_per_position.html', context)

def announce_winners(request, entity):
    current_year = SACYear.objects.get_current()
    per_city = []
    for city_code, city_name in city_choices:
        if not current_year.is_city_results_due(city_code):
            continue
        positions = Position.objects.filter(entity=entity,
                                            city=city_code,
                                            year=current_year)
        per_city.append((city_name, positions))
    context = {'entity': entity, 'per_city': per_city}
    return render(request, 'voting/announce_winners.html', context)

class UserAutocomplete(autocomplete.Select2QuerySetView):
    def get_queryset(self):
        if not self.request.user.is_authenticated():
            return User.objects.none()

        qs = User.objects.filter(is_active=True)

        if self.q:
            search_fields = [re.sub('^user__', '', field) for field in utils.BASIC_SEARCH_FIELDS]
            print search_fields
            qs = utils.get_search_queryset(qs, search_fields, self.q)

        return qs

    def get_result_label(self, item):
        full_name = accounts.utils.get_user_ar_full_name(item) or item.username
        return"%s <bdi style='font-family: monospace;'>(%s)</bdi>" % (full_name, item.username)
<commit_msg>add year filter to ticket<commit_after># -*- coding: utf-8  -*-
from dal import autocomplete
from django.core.exceptions import PermissionDenied
from django.core.urlresolvers import reverse
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.db.models import Count
from django.http import Http404, HttpResponseRedirect
from django.shortcuts import render, get_object_or_404
from django.views.decorators import csrf
from .forms import NominationForm
from .models import SACYear, Position, Nomination, \
                    NominationAnnouncement, VoteNomination, \
                    city_choices
from . import decorators, utils
import accounts.utils
import re

def show_index(request):
    if request.user.is_authenticated():
        return render(request, 'accounts/home.html')
    else:
        return render(request, 'accounts/home_unauthenticated.html')

@login_required
def list_positions(request, entity):
    current_year = SACYear.objects.get_current()
    if not current_year.is_nomination_open():
        return render(request, "voting/closed.html")

    user_nominations = Nomination.objects.filter(position__entity=entity,
                                                 position__year=current_year,
                                                 user=request.user)
    positions = Position.objects.filter(entity=entity, year=current_year)
    if not request.user.is_superuser:
        positions = positions.filter(colleges_allowed_to_nominate=request.user.profile.college)

    context = {'user_nominations': user_nominations,
               'positions': positions}
    return render(request, 'voting/list_' + entity + '_positions.html', context)

@login_required
def add_nominee(request, position_id):
    current_year = SACYear.objects.get_current()
    if not current_year.is_nomination_open():
        return HttpResponseRedirect(reverse("voting:closed"))

    position = get_object_or_404(Position, pk=position_id)
    context = {'position': position}

    user_nominations = Nomination.objects.filter(position__entity=position.entity,
                                                 position__year=current_year,
                                                 user=request.user)
    if user_nominations.exists():
        context['already_on'] = True
    else:
        if request.user.is_superuser:
            positions = Position.objects.filter(year=current_year)
        else:
            positions = Position.objects.filter(colleges_allowed_to_nominate=request.user.profile.college,
                                                year=current_year)
            if position not in positions:
                raise PermissionDenied
        if request.method == 'POST':
            instance = Nomination(user=request.user, position=position)
            form = NominationForm(request.POST, request.FILES, instance=instance)
            if form.is_valid():
                instance = form.save()
                return HttpResponseRedirect(reverse("voting:nomination_thanks", args=(position.pk,)))
        elif request.method == 'GET':
            form = NominationForm()
        context['form'] = form

    return render(request,'voting/add_nominee.html', context)

@login_required
def show_nomination(request, position_id, nomination_id):
    nomination = get_object_or_404(Nomination, position__pk=position_id,
                                   pk=nomination_id)
    if not request.user.is_superuser and \
       not nomination.user == request.user:
        raise PermissionDenied
    context = {'nomination': nomination}
    return render(request,'voting/show_nomination.html', context)

def announce_nominees(request, entity):
    current_year = SACYear.objects.get_current()
    context = {}
    if current_year.is_announcement_due():
        per_city = []
        for city_code, city_name in city_choices:
            positions = Position.objects\
                               .filter(entity=entity, city=city_code, year=current_year)\
                               .annotate(nomination_count=Count('nominationannouncement'))\
                               .filter(nomination_count__gt=1)
            if positions.exists():
                per_city.append((city_name, positions))
        context = {'entity': entity, 'per_city': per_city}

    return render(request, 'voting/announce_nominees.html', context)

@login_required
def show_voting_index(request):
    current_year = SACYear.objects.get_current()
    if not request.user.is_superuser and not current_year.is_voting_open():
        return render(request, "voting/voting_closed.html")
    else:
        position_pool = Position.objects.filter(year=current_year)\
                                        .annotate(announced_count=Count('nominationannouncement'))\
                                        .filter(announced_count__gte=2)\
                                        .exclude(votenomination__user=request.user)\
                                        .order_by('entity')
        if request.user.profile:
            position_pool = position_pool.filter(colleges_allowed_to_vote=request.user.profile.college)

        if position_pool.exists():
            completed_voting = False
            votes = VoteNomination.objects.none()
        else:
            completed_voting = True
            votes = VoteNomination.objects.filter(user=request.user,position__year=current_year)
        qrcode_value = utils.get_ticket(request.user)
        context = {'qrcode_value': qrcode_value,
                   'completed_voting': completed_voting,
                   'votes': votes}
        return render(request,'voting/show_voting.html', context)

@login_required
@decorators.ajax_only
@decorators.post_only
@csrf.csrf_exempt
def handle_vote(request):
    current_year = SACYear.objects.get_current()
    if not request.user.is_superuser and not current_year.is_voting_open():
        raise Exception("التصويت غير مُتاح حاليًا!")

    nomination_vote_pk = request.POST.get('nomination_vote_pk', None)
    if nomination_vote_pk:
        if nomination_vote_pk == 'skip':
            position_pk = request.POST.get('position_pk')
            nomination_announcement = None
            position = Position.objects.get(pk=position_pk)
        else:
            nomination_announcement = NominationAnnouncement.objects.get(pk=nomination_vote_pk)
            position = nomination_announcement.position

        previous_vote = VoteNomination.objects.filter(position=position,
                                                      user=request.user).exists()
        if previous_vote:
            raise Exception(u'سبق أن صوتّ لهذا المنصب')

        if not request.user.is_superuser and \
           not request.user.profile.college in\
               position.colleges_allowed_to_vote.all():
            raise PermissionDenied
        else:
            VoteNomination.objects.create(nomination_announcement=nomination_announcement,
                                          position=position,
                                          user=request.user)

    position_pool = Position.objects.filter(year=current_year)\
                                    .annotate(announced_count=Count('nominationannouncement'))\
                                    .filter(announced_count__gte=2)\
                                    .exclude(votenomination__user=request.user)
    if request.user.is_superuser:
        next_position = position_pool.first()
    else:
        next_position = position_pool.filter(colleges_allowed_to_vote=request.user.profile.college)\
                                     .first()
    if next_position:
        nominations = []
        for nomination in next_position.nominationannouncement_set.order_by('user__profile__ar_first_name'):
            nomination = {'pk': nomination.pk,
                          'nominee_name': nomination.user.profile.get_ar_full_name(),
                          'plan': nomination.plan.url,
                          'cv': nomination.cv.url}
            nominations.append(nomination)
        return {"position_title": next_position.title,
                "position_pk": next_position.pk,
                "entity": next_position.entity,
                "note": next_position.note,
                "nominations": nominations}
    else:
        return {'done': 1}

@login_required
def indicators(request):
    if not request.user.is_superuser:
        raise PermissionDenied
    current_year = SACYear.objects.get_current()
    election_positions = Position.objects.filter(year=current_year)\
                                         .annotate(announced_count=Count('nominationannouncement'))\
                                         .filter(announced_count__gte=2).order_by('entity', 'city')
    context = {'election_positions': election_positions}

    return render(request, 'voting/indicators.html', context)

@login_required
def list_votes_per_position(request, position_id):
    if not request.user.is_superuser:
        raise PermissionDenied

    position = get_object_or_404(Position, pk=position_id)
    votes = VoteNomination.objects.filter(position=position)
    context = {'position': position, 'votes': votes}

    return render(request, 'voting/list_votes_per_position.html', context)

def announce_winners(request, entity):
    current_year = SACYear.objects.get_current()
    per_city = []
    for city_code, city_name in city_choices:
        if not current_year.is_city_results_due(city_code):
            continue
        positions = Position.objects.filter(entity=entity,
                                            city=city_code,
                                            year=current_year)
        per_city.append((city_name, positions))
    context = {'entity': entity, 'per_city': per_city}
    return render(request, 'voting/announce_winners.html', context)

class UserAutocomplete(autocomplete.Select2QuerySetView):
    def get_queryset(self):
        if not self.request.user.is_authenticated():
            return User.objects.none()

        qs = User.objects.filter(is_active=True)

        if self.q:
            search_fields = [re.sub('^user__', '', field) for field in utils.BASIC_SEARCH_FIELDS]
            print search_fields
            qs = utils.get_search_queryset(qs, search_fields, self.q)

        return qs

    def get_result_label(self, item):
        full_name = accounts.utils.get_user_ar_full_name(item) or item.username
        return"%s <bdi style='font-family: monospace;'>(%s)</bdi>" % (full_name, item.username)
<|endoftext|>


<commit_before>import struct
from binascii import crc32
from io import FileIO
from io import open as fopen
import os

__version__ = "0.11"
__author__ = "Rossen Georgiev"

def open(vpk_path, **kwargs):
    """
    Returns a VPK instance for specified path. Same argumets as VPK class.
    """
    return VPK(vpk_path, **kwargs)

def new(dir_path):
    """
    Returns a NewVPK instance for the specific path.
    """
    return NewVPK(dir_path)

class NewVPK:
    def __init__(self, path):
        self.signature = 0x55aa1234
        self.version = 1
        self.tree_length = 0
        self.header_length = 4*3

        self.tree = {}
        self.path = ''
        self.file_count = 0

        self.read_dir(path)

    def __repr__(self):
        return "%s('%s')" % (self.__class__.__name__, self.path)

    def read_dir(self, path):
        """
        Reads the given path into the tree
        """
        self.tree = {}
        self.file_count = 0
        self.path = path

        for root, _, filelist in os.walk(path):
            rel = root[len(path)+1:]

            # empty rel, means file is in root dir
            if not rel:
                rel = ' '

            for filename in filelist:
                filename = filename.split('.')
                if len(filename) <= 1:
                    raise RuntimeError("Files without an extension are not supported: {0}".format(
                                       repr(os.path.join(root, '.'.join(filename))),
                                       ))

                ext = filename[-1]
                filename = '.'.join(filename[:-1])

                if ext not in self.tree:
                    self.tree[ext] = {}
                if rel not in self.tree[ext]:
                    self.tree[ext][rel] = []

                self.tree[ext][rel].append(filename)
                self.file_count += 1

        self.tree_length = self.calculate_tree_length()

    def calculate_tree_length(self):
        """
        Walks the tree and calculate the tree length
        """
        tree_length = 0

        for ext in self.tree:
            tree_length += len(ext) + 2

            for relpath in self.tree[ext]:
                tree_length += len(relpath) + 2

                for filename in self.tree[ext][relpath]:
                    tree_length += len(filename) + 1 + 18

        return tree_length + 1

    def save(self, vpk_output_path):
        """
        Saves the VPK at the given path
        """
        with fopen(vpk_output_path, 'wb') as f:
            # write VPK1 header
            f.write(struct.pack("3I", self.signature,
                                      self.version,
                                      self.tree_length))

            self.header_length = f.tell()

            data_offset = self.header_length + self.tree_length

            # write file tree
            for ext in self.tree:
                f.write("{0}\x00".format(ext).encode('latin-1'))

                for relpath in self.tree[ext]:
                    f.write("{0}\x00".format(relpath).encode('latin-1'))

                    for filename in self.tree[ext][relpath]:
                        f.write("{0}\x00".format(filename).encode('latin-1'))

                        # append file data
                        metadata_offset = f.tell()
                        file_offset = data_offset
                        real_filename = filename if not ext else "{0}.{1}".format(filename, ext)
                        checksum = 0
                        f.seek(data_offset)

                        with fopen(os.path.join(self.path,
                                                '' if relpath == ' ' else relpath,
                                                real_filename
                                                ),
                                   'rb') as pakfile:
                            for chunk in iter(lambda: pakfile.read(1024), b''):
                                checksum = crc32(chunk, checksum)
                                f.write(chunk)

                        data_offset = f.tell()
                        file_length = f.tell() - file_offset
                        f.seek(metadata_offset)

                        # metadata

                        # crc32
                        # preload_length
                        # archive_index
                        # archive_offset
                        # file_length
                        # term
                        f.write(struct.pack("IHHIIH", checksum & 0xFFffFFff,
                                                      0,
                                                      0x7fff,
                                                      file_offset - self.tree_length - self.header_length,
                                                      file_length,
                                                      0xffff
                                                      ))

                    # next relpath
                    f.write(b"\x00")
                # next ext
                f.write(b"\x00")
            # end of file tree
            f.write(b"\x00")

    def save_and_open(self, path):
        """
        Saves the VPK file and returns VPK instance of it
        """
        self.save(path)
        return VPK(path)

class VPK:
    """
    Wrapper for reading Valve's VPK files
    """

    def __init__(self, vpk_path, read_header_only=False):
        # header
        self.signature = 0
        self.version = 0
        self.tree_length = 0
        self.header_length = 0

        self.tree = {}
        self.vpk_path = vpk_path

        self.read_header()

        if not read_header_only:
            self.read_index()

    def __repr__(self):
        headonly = ', read_header_only=True' if len(self) == 0 else ''
        return "%s('%s'%s)" % (self.__class__.__name__, self.vpk_path, headonly)

    def __iter__(self):
        return self.tree.__iter__()

    def items(self):
        def items_generator(tree):
            for path in tree:
                yield path, self.get_file_meta(path)

        return items_generator(self.tree)

    def __len__(self):
        return len(self.tree)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pass

    def _read_sz(self, f):
        buf = b''

        for chunk in iter(lambda: f.read(64), b''):
            pos = chunk.find(b'\x00')
            if pos > -1:
                buf += chunk[:pos]
                f.seek(f.tell() - (len(chunk) - (pos + 1)))
                break

            buf += chunk

        return buf.decode('ascii')

    def __getitem__(self, key):
        """
        Returns VPKFile instance
        """
        return self.get_file(key)

    def get_file(self, path):
        """
        Returns VPKFile instance for the given path
        """
        metadata = self.get_file_meta(path)
        return VPKFile(self.vpk_path, filepath=path, **metadata)

    def get_file_meta(self, path):
        """
        Returns metadata for given file path
        """
        if path not in self.tree:
            raise KeyError("Path doesn't exist")

        return dict(zip(['preload',
                         'crc32',
                         'preload_length',
                         'archive_index',
                         'archive_offset',
                         'file_length',
                         ], self.tree[path]))

    def read_header(self):
        """
        Reads VPK file header from the file
        """
        with fopen(self.vpk_path, 'rb') as f:
            (self.signature,
             self.version,
             self.tree_length
             ) = struct.unpack("3I", f.read(3*4))

            # original format - headerless
            if self.signature != 0x55aa1234:
                self.signature = 0
                self.version = 0
                self.tree_length = 0
            # v1
            elif self.version == 1:
                self.header_length += 4*3
            # v2 with extended header
            #
            # according to http://forum.xentax.com/viewtopic.php?f=10&t=11208
            # struct VPKDirHeader_t
            # {
            #    int32 m_nHeaderMarker;
            #    int32 m_nVersion;
            #    int32 m_nDirectorySize;
            #    int32 m_nEmbeddedChunkSize;
            #    int32 m_nChunkHashesSize;
            #    int32 m_nSelfHashesSize;
            #    int32 m_nSignatureSize;
            # }
            elif self.version == 2:
                (self.embed_chunk_length,
                 self.chunk_hashes_length,
                 self.self_hashes_length,
                 self.signature_length
                 ) = struct.unpack("4I", f.read(4*4))
                self.header_length += 4*7
            else:
                raise ValueError("Invalid header, or unsupported version")

    def read_index(self):
        """
        Reads the index and populates the directory tree
        """

        self.tree = {}
        with fopen(self.vpk_path, 'rb') as f:
            f.seek(self.header_length)

            while True:
                if self.version > 0 and f.tell() > self.tree_length + self.header_length:
                    raise ValueError("Error parsing index (out of bounds)")

                ext = self._read_sz(f)
                if ext == '':
                    break

                while True:
                    path = self._read_sz(f)
                    if path == '':
                        break
                    if path != ' ':
                        path += '/'
                    else:
                        path = ''

                    while True:
                        name = self._read_sz(f)
                        if name == '':
                            break

                        # crc32
                        # preload_length
                        # archive_index
                        # archive_offset
                        # file_length
                        metadata = list(struct.unpack("IHHII", f.read(16)))

                        if struct.unpack("H", f.read(2))[0] != 0xffff:
                            raise ValueError("Error while parsing index")

                        if metadata[2] == 0x7fff:
                            metadata[3] += self.header_length + self.tree_length

                        metadata.insert(0, f.read(metadata[1]))

                        self.tree["{0}{1}.{2}".format(path, name, ext)] = tuple(metadata)

class VPKFile(FileIO):
    """
    Wrapper class for files with VPK

    Should act like a regular file object. No garantees
    """

    def __init__(self, vpk_path, **kw):
        self.vpk_path = vpk_path
        self.vpk_meta = kw

        for k, v in kw.items():
            setattr(self, k, v)

        if self.vpk_meta['preload'] != b'':
            self.vpk_meta['preload'] = '...'

        # total file length
        self.length = self.preload_length + self.file_length
        # offset of entire file
        self.offset = 0

        if self.file_length == 0:
            self.vpk_path = None
            return

        super(VPKFile, self).__init__(vpk_path.replace("dir.", "%03d." % self.archive_index), 'rb')
        super(VPKFile, self).seek(self.archive_offset)

    def save(self, path):
        """
        Save the file to the specified path
        """
        # remember and restore file position
        pos = self.tell()
        self.seek(0)

        with fopen(path, 'wb') as output:
            output.truncate(self.length)
            for chunk in iter(lambda: self.read(1024), b''):
                output.write(chunk)

        self.seek(pos)

    def verify(self):
        """
        Returns True if the file contents match with the CRC32 attribute

        note: reset
        """

        # remember file pointer
        pos = self.tell()
        self.seek(0)

        checksum = 0
        for chunk in iter(lambda: self.read(1024), b''):
            checksum = crc32(chunk, checksum)

        # restore file pointer
        self.seek(pos)

        return self.crc32 == checksum & 0xffffffff

    def __repr__(self):
        return "%s(%s, %s)" % (
            self.__class__.__name__,
            repr(self.name) if self.file_length > 0 else None,
            ', '.join(["%s=%s" % (k, repr(v)) for k, v in self.vpk_meta.items()])
            )

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        line = self.readline()
        if line == b'':
            raise StopIteration
        return line

    def close(self):
        super(VPKFile, self).close()

    def tell(self):
        return self.offset

    def seek(self, offset):
        if offset < 0:
            raise IOError("Invalid argument")

        self.offset = offset
        if self.file_length > 0:
            super(VPKFile, self).seek(self.archive_offset + max(offset - self.preload_length, 0))

    def readlines(self):
        return [line for line in self]

    def readline(self, a=False):
        buf = b''

        for chunk in iter(lambda: self.read(256), b''):
            pos = chunk.find(b'\n')
            if pos > -1:
                pos += 1  # include \n
                buf += chunk[:pos]
                self.seek(self.offset - (len(chunk) - pos))
                break

            buf += chunk

        return buf

    def read(self, length=-1):
        if length == 0 or self.offset >= self.length:
            return b''

        data = b''

        if self.offset <= self.preload_length:
            data += self.preload[self.offset:self.offset+length if length > -1 else None]
            self.offset += len(data)
            if length > 0:
                length = max(length - len(data), 0)

        if self.file_length > 0 and self.offset >= self.preload_length:
            left = self.file_length - (self.offset - self.preload_length)
            data += super(VPKFile, self).read(left if length == -1 else min(left, length))
            self.offset += left if length == -1 else min(left, length)

        return data

    def write(self, seq):
        raise NotImplementedError("write method is not supported")
<commit_msg>fix bug when forming the relative path<commit_after>import struct
from binascii import crc32
from io import FileIO
from io import open as fopen
import os

__version__ = "0.11"
__author__ = "Rossen Georgiev"

def open(vpk_path, **kwargs):
    """
    Returns a VPK instance for specified path. Same argumets as VPK class.
    """
    return VPK(vpk_path, **kwargs)

def new(dir_path):
    """
    Returns a NewVPK instance for the specific path.
    """
    return NewVPK(dir_path)

class NewVPK:
    def __init__(self, path):
        self.signature = 0x55aa1234
        self.version = 1
        self.tree_length = 0
        self.header_length = 4*3

        self.tree = {}
        self.path = ''
        self.file_count = 0

        self.read_dir(path)

    def __repr__(self):
        return "%s('%s')" % (self.__class__.__name__, self.path)

    def read_dir(self, path):
        """
        Reads the given path into the tree
        """
        self.tree = {}
        self.file_count = 0
        self.path = path

        for root, _, filelist in os.walk(path):
            rel = root[len(path):].lstrip('/')

            # empty rel, means file is in root dir
            if not rel:
                rel = ' '

            for filename in filelist:
                filename = filename.split('.')
                if len(filename) <= 1:
                    raise RuntimeError("Files without an extension are not supported: {0}".format(
                                       repr(os.path.join(root, '.'.join(filename))),
                                       ))

                ext = filename[-1]
                filename = '.'.join(filename[:-1])

                if ext not in self.tree:
                    self.tree[ext] = {}
                if rel not in self.tree[ext]:
                    self.tree[ext][rel] = []

                self.tree[ext][rel].append(filename)
                self.file_count += 1

        self.tree_length = self.calculate_tree_length()

    def calculate_tree_length(self):
        """
        Walks the tree and calculate the tree length
        """
        tree_length = 0

        for ext in self.tree:
            tree_length += len(ext) + 2

            for relpath in self.tree[ext]:
                tree_length += len(relpath) + 2

                for filename in self.tree[ext][relpath]:
                    tree_length += len(filename) + 1 + 18

        return tree_length + 1

    def save(self, vpk_output_path):
        """
        Saves the VPK at the given path
        """
        with fopen(vpk_output_path, 'wb') as f:
            # write VPK1 header
            f.write(struct.pack("3I", self.signature,
                                      self.version,
                                      self.tree_length))

            self.header_length = f.tell()

            data_offset = self.header_length + self.tree_length

            # write file tree
            for ext in self.tree:
                f.write("{0}\x00".format(ext).encode('latin-1'))

                for relpath in self.tree[ext]:
                    f.write("{0}\x00".format(relpath).encode('latin-1'))

                    for filename in self.tree[ext][relpath]:
                        f.write("{0}\x00".format(filename).encode('latin-1'))

                        # append file data
                        metadata_offset = f.tell()
                        file_offset = data_offset
                        real_filename = filename if not ext else "{0}.{1}".format(filename, ext)
                        checksum = 0
                        f.seek(data_offset)

                        with fopen(os.path.join(self.path,
                                                '' if relpath == ' ' else relpath,
                                                real_filename
                                                ),
                                   'rb') as pakfile:
                            for chunk in iter(lambda: pakfile.read(1024), b''):
                                checksum = crc32(chunk, checksum)
                                f.write(chunk)

                        data_offset = f.tell()
                        file_length = f.tell() - file_offset
                        f.seek(metadata_offset)

                        # metadata

                        # crc32
                        # preload_length
                        # archive_index
                        # archive_offset
                        # file_length
                        # term
                        f.write(struct.pack("IHHIIH", checksum & 0xFFffFFff,
                                                      0,
                                                      0x7fff,
                                                      file_offset - self.tree_length - self.header_length,
                                                      file_length,
                                                      0xffff
                                                      ))

                    # next relpath
                    f.write(b"\x00")
                # next ext
                f.write(b"\x00")
            # end of file tree
            f.write(b"\x00")

    def save_and_open(self, path):
        """
        Saves the VPK file and returns VPK instance of it
        """
        self.save(path)
        return VPK(path)

class VPK:
    """
    Wrapper for reading Valve's VPK files
    """

    def __init__(self, vpk_path, read_header_only=False):
        # header
        self.signature = 0
        self.version = 0
        self.tree_length = 0
        self.header_length = 0

        self.tree = {}
        self.vpk_path = vpk_path

        self.read_header()

        if not read_header_only:
            self.read_index()

    def __repr__(self):
        headonly = ', read_header_only=True' if len(self) == 0 else ''
        return "%s('%s'%s)" % (self.__class__.__name__, self.vpk_path, headonly)

    def __iter__(self):
        return self.tree.__iter__()

    def items(self):
        def items_generator(tree):
            for path in tree:
                yield path, self.get_file_meta(path)

        return items_generator(self.tree)

    def __len__(self):
        return len(self.tree)

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pass

    def _read_sz(self, f):
        buf = b''

        for chunk in iter(lambda: f.read(64), b''):
            pos = chunk.find(b'\x00')
            if pos > -1:
                buf += chunk[:pos]
                f.seek(f.tell() - (len(chunk) - (pos + 1)))
                break

            buf += chunk

        return buf.decode('ascii')

    def __getitem__(self, key):
        """
        Returns VPKFile instance
        """
        return self.get_file(key)

    def get_file(self, path):
        """
        Returns VPKFile instance for the given path
        """
        metadata = self.get_file_meta(path)
        return VPKFile(self.vpk_path, filepath=path, **metadata)

    def get_file_meta(self, path):
        """
        Returns metadata for given file path
        """
        if path not in self.tree:
            raise KeyError("Path doesn't exist")

        return dict(zip(['preload',
                         'crc32',
                         'preload_length',
                         'archive_index',
                         'archive_offset',
                         'file_length',
                         ], self.tree[path]))

    def read_header(self):
        """
        Reads VPK file header from the file
        """
        with fopen(self.vpk_path, 'rb') as f:
            (self.signature,
             self.version,
             self.tree_length
             ) = struct.unpack("3I", f.read(3*4))

            # original format - headerless
            if self.signature != 0x55aa1234:
                self.signature = 0
                self.version = 0
                self.tree_length = 0
            # v1
            elif self.version == 1:
                self.header_length += 4*3
            # v2 with extended header
            #
            # according to http://forum.xentax.com/viewtopic.php?f=10&t=11208
            # struct VPKDirHeader_t
            # {
            #    int32 m_nHeaderMarker;
            #    int32 m_nVersion;
            #    int32 m_nDirectorySize;
            #    int32 m_nEmbeddedChunkSize;
            #    int32 m_nChunkHashesSize;
            #    int32 m_nSelfHashesSize;
            #    int32 m_nSignatureSize;
            # }
            elif self.version == 2:
                (self.embed_chunk_length,
                 self.chunk_hashes_length,
                 self.self_hashes_length,
                 self.signature_length
                 ) = struct.unpack("4I", f.read(4*4))
                self.header_length += 4*7
            else:
                raise ValueError("Invalid header, or unsupported version")

    def read_index(self):
        """
        Reads the index and populates the directory tree
        """

        self.tree = {}
        with fopen(self.vpk_path, 'rb') as f:
            f.seek(self.header_length)

            while True:
                if self.version > 0 and f.tell() > self.tree_length + self.header_length:
                    raise ValueError("Error parsing index (out of bounds)")

                ext = self._read_sz(f)
                if ext == '':
                    break

                while True:
                    path = self._read_sz(f)
                    if path == '':
                        break
                    if path != ' ':
                        path += '/'
                    else:
                        path = ''

                    while True:
                        name = self._read_sz(f)
                        if name == '':
                            break

                        # crc32
                        # preload_length
                        # archive_index
                        # archive_offset
                        # file_length
                        metadata = list(struct.unpack("IHHII", f.read(16)))

                        if struct.unpack("H", f.read(2))[0] != 0xffff:
                            raise ValueError("Error while parsing index")

                        if metadata[2] == 0x7fff:
                            metadata[3] += self.header_length + self.tree_length

                        metadata.insert(0, f.read(metadata[1]))

                        self.tree["{0}{1}.{2}".format(path, name, ext)] = tuple(metadata)

class VPKFile(FileIO):
    """
    Wrapper class for files with VPK

    Should act like a regular file object. No garantees
    """

    def __init__(self, vpk_path, **kw):
        self.vpk_path = vpk_path
        self.vpk_meta = kw

        for k, v in kw.items():
            setattr(self, k, v)

        if self.vpk_meta['preload'] != b'':
            self.vpk_meta['preload'] = '...'

        # total file length
        self.length = self.preload_length + self.file_length
        # offset of entire file
        self.offset = 0

        if self.file_length == 0:
            self.vpk_path = None
            return

        super(VPKFile, self).__init__(vpk_path.replace("dir.", "%03d." % self.archive_index), 'rb')
        super(VPKFile, self).seek(self.archive_offset)

    def save(self, path):
        """
        Save the file to the specified path
        """
        # remember and restore file position
        pos = self.tell()
        self.seek(0)

        with fopen(path, 'wb') as output:
            output.truncate(self.length)
            for chunk in iter(lambda: self.read(1024), b''):
                output.write(chunk)

        self.seek(pos)

    def verify(self):
        """
        Returns True if the file contents match with the CRC32 attribute

        note: reset
        """

        # remember file pointer
        pos = self.tell()
        self.seek(0)

        checksum = 0
        for chunk in iter(lambda: self.read(1024), b''):
            checksum = crc32(chunk, checksum)

        # restore file pointer
        self.seek(pos)

        return self.crc32 == checksum & 0xffffffff

    def __repr__(self):
        return "%s(%s, %s)" % (
            self.__class__.__name__,
            repr(self.name) if self.file_length > 0 else None,
            ', '.join(["%s=%s" % (k, repr(v)) for k, v in self.vpk_meta.items()])
            )

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        line = self.readline()
        if line == b'':
            raise StopIteration
        return line

    def close(self):
        super(VPKFile, self).close()

    def tell(self):
        return self.offset

    def seek(self, offset):
        if offset < 0:
            raise IOError("Invalid argument")

        self.offset = offset
        if self.file_length > 0:
            super(VPKFile, self).seek(self.archive_offset + max(offset - self.preload_length, 0))

    def readlines(self):
        return [line for line in self]

    def readline(self, a=False):
        buf = b''

        for chunk in iter(lambda: self.read(256), b''):
            pos = chunk.find(b'\n')
            if pos > -1:
                pos += 1  # include \n
                buf += chunk[:pos]
                self.seek(self.offset - (len(chunk) - pos))
                break

            buf += chunk

        return buf

    def read(self, length=-1):
        if length == 0 or self.offset >= self.length:
            return b''

        data = b''

        if self.offset <= self.preload_length:
            data += self.preload[self.offset:self.offset+length if length > -1 else None]
            self.offset += len(data)
            if length > 0:
                length = max(length - len(data), 0)

        if self.file_length > 0 and self.offset >= self.preload_length:
            left = self.file_length - (self.offset - self.preload_length)
            data += super(VPKFile, self).read(left if length == -1 else min(left, length))
            self.offset += left if length == -1 else min(left, length)

        return data

    def write(self, seq):
        raise NotImplementedError("write method is not supported")
<|endoftext|>


<commit_before>import numpy as np
from scipy import sparse

from vsm.model import BaseModel

class TfModel(BaseModel):
    """
    """
    def __init__(self, corpus, tok_name):
        """
        """
        self.corpus = corpus.corpus
        self.docs = corpus.view_tokens(tok_name)
        self.shape = (corpus.terms.size, len(self.docs))

    def train(self):
        """
        """
        print 'Computing term frequencies'
        data = np.ones_like(self.corpus)
        row_indices = self.corpus
        col_indices = np.empty_like(self.corpus)
        j, k = 0, 0

        for i,token in enumerate(self.docs):
            k += len(token)
            col_indices[j:k] = i
            j = k

        coo_in = (data, (row_indices, col_indices))
        self.matrix = sparse.coo_matrix(coo_in, shape=self.shape,
                                        dtype=np.int32)

def test_TfModel():

    from vsm.util.corpustools import random_corpus

    c = random_corpus(1000, 100, 0, 20, tok_name='document')

    m = TfModel(c, 'document')
    m.train()

    from tempfile import NamedTemporaryFile
    import os

    try:
        tmp = NamedTemporaryFile(delete=False, suffix='.npz')
        m.save(tmp.name)
        tmp.close()
        m1 = TfModel.load(tmp.name)
        assert (m.matrix.todense() == m1.matrix.todense()).all()

    finally:
        os.remove(tmp.name)

    return m.matrix
<commit_msg>TfModel instances now remember their document names<commit_after>import numpy as np
from scipy import sparse

from vsm.model import BaseModel

class TfModel(BaseModel):
    """
    """
    def __init__(self, corpus, tok_name):
        """
        """
        self.corpus = corpus
        self.tok_name = tok_name

    def train(self):
        """
        """
        docs = self.corpus.view_tokens(self.tok_name)
        shape = (self.corpus.terms.size, len(docs))

        print 'Computing term frequencies'
        data = np.ones_like(self.corpus.corpus)
        row_indices = self.corpus.corpus
        col_indices = np.empty_like(self.corpus.corpus)
        j, k = 0, 0

        for i,token in enumerate(docs):
            k += len(token)
            col_indices[j:k] = i
            j = k

        coo_in = (data, (row_indices, col_indices))
        self.matrix = sparse.coo_matrix(coo_in, shape=shape,
                                        dtype=np.int32)

def test_TfModel():

    from vsm.util.corpustools import random_corpus

    c = random_corpus(10000, 100, 0, 100, tok_name='document')

    m = TfModel(c, 'document')
    m.train()

    from tempfile import NamedTemporaryFile
    import os

    try:
        tmp = NamedTemporaryFile(delete=False, suffix='.npz')
        m.save(tmp.name)
        tmp.close()
        m1 = TfModel.load(tmp.name)
        assert (m.matrix.todense() == m1.matrix.todense()).all()

    finally:
        os.remove(tmp.name)

    return m.matrix
<|endoftext|>


<commit_before># -*- test-case-name: vumi.tests.test_message -*-

import json
from uuid import uuid4
from datetime import datetime

from errors import MissingMessageField, InvalidMessageField

from vumi.utils import to_kwargs

# This is the date format we work with internally
VUMI_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f"

def date_time_decoder(json_object):
    for key, value in json_object.items():
        try:
            json_object[key] = datetime.strptime(value,
                    VUMI_DATE_FORMAT)
        except ValueError:
            continue
        except TypeError:
            continue
    return json_object

class JSONMessageEncoder(json.JSONEncoder):
    """A JSON encoder that is able to serialize datetime"""
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime(VUMI_DATE_FORMAT)
        return super(JSONMessageEncoder, self).default(obj)

def from_json(json_string):
    return json.loads(json_string, object_hook=date_time_decoder)

def to_json(obj):
    return json.dumps(obj, cls=JSONMessageEncoder)

class Message(object):
    """
    Start of a somewhat unified message object to be
    used internally in Vumi and while being in transit
    over AMQP

    scary transport format -> Vumi Tansport -> Unified Message -> Vumi Worker

    """

    def __init__(self, _process_fields=True, **kwargs):
        if _process_fields:
            kwargs = self.process_fields(kwargs)
        self.payload = kwargs
        self.validate_fields()

    def process_fields(self, fields):
        return fields

    def validate_fields(self):
        pass

    def assert_field_present(self, *fields):
        for field in fields:
            if field not in self.payload:
                raise MissingMessageField(field)

    def assert_field_value(self, field, *values):
        self.assert_field_present(field)
        if self.payload[field] not in values:
            raise InvalidMessageField(field)

    def to_json(self):
        return to_json(self.payload)

    @classmethod
    def from_json(cls, json_string):
        return cls(_process_fields=False, **to_kwargs(from_json(json_string)))

    def __str__(self):
        return u"<Message payload=\"%s\">" % repr(self.payload)

    def __repr__(self):
        return str(self)

    def __eq__(self, other):
        if isinstance(other, Message):
            return self.payload == other.payload
        return False

    def __contains__(self, key):
        return key in self.payload

    def __getitem__(self, key):
        return self.payload[key]

    def __setitem__(self, key, value):
        self.payload[key] = value

    def get(self, key, default=None):
        return self.payload.get(key, default)

    def items(self):
        return self.payload.items()

    def copy(self):
        return self.from_json(self.to_json())

class TransportMessage(Message):
    """Common base class for messages sent to or from a transport."""

    # sub-classes should set the message type
    MESSAGE_TYPE = None
    MESSAGE_VERSION = '20110921'
    DEFAULT_ENDPOINT_NAME = 'default'

    @staticmethod
    def generate_id():
        """
        Generate a unique message id.

        There are places where we want a message id before we can
        build a complete message. This lets us do that in a consistent
        manner.
        """
        return uuid4().get_hex()

    def process_fields(self, fields):
        fields.setdefault('message_version', self.MESSAGE_VERSION)
        fields.setdefault('message_type', self.MESSAGE_TYPE)
        fields.setdefault('timestamp', datetime.utcnow())
        fields.setdefault('routing_metadata', {})
        fields.setdefault('helper_metadata', {})
        return fields

    def validate_fields(self):
        self.assert_field_value('message_version', self.MESSAGE_VERSION)
        # We might get older event messages without the `helper_metadata`
        # field.
        self.payload.setdefault('helper_metadata', {})
        self.assert_field_present(
            'message_type',
            'timestamp',
            'helper_metadata',
            )
        if self['message_type'] is None:
            raise InvalidMessageField('message_type')

    @property
    def routing_metadata(self):
        return self.payload.setdefault('routing_metadata', {})

    @classmethod
    def check_routing_endpoint(cls, endpoint_name):
        if endpoint_name is None:
            return cls.DEFAULT_ENDPOINT_NAME
        return endpoint_name

    def set_routing_endpoint(self, endpoint_name=None):
        endpoint_name = self.check_routing_endpoint(endpoint_name)
        self.routing_metadata['endpoint_name'] = endpoint_name

    def get_routing_endpoint(self):
        endpoint_name = self.routing_metadata.get('endpoint_name')
        return self.check_routing_endpoint(endpoint_name)

class TransportUserMessage(TransportMessage):
    """Message to or from a user.

    transport_type = sms, ussd, etc
    helper_metadata = for use by dispathers and off-to-the-side
                      components like failure workers (not for use
                      by transports or message workers).
    """

    MESSAGE_TYPE = 'user_message'

    # session event constants
    #
    # SESSION_NONE, SESSION_NEW, SESSION_RESUME, and SESSION_CLOSE
    # may be sent from the transport to a worker. SESSION_NONE indicates
    # there is no relevant session for this message.
    #
    # SESSION_NONE and SESSION_CLOSE may be sent from the worker to
    # the transport. SESSION_NONE indicates any existing session
    # should be continued. SESSION_CLOSE indicates that any existing
    # session should be terminated after sending the message.
    SESSION_NONE, SESSION_NEW, SESSION_RESUME, SESSION_CLOSE = (
        None, 'new', 'resume', 'close')

    # list of valid session events
    SESSION_EVENTS = frozenset([SESSION_NONE, SESSION_NEW, SESSION_RESUME,
                                SESSION_CLOSE])

    # canonical transport types
    TT_HTTP_API = 'http_api'
    TT_IRC = 'irc'
    TT_TELNET = 'telnet'
    TT_TWITTER = 'twitter'
    TT_SMS = 'sms'
    TT_USSD = 'ussd'
    TT_XMPP = 'xmpp'
    TT_MXIT = 'mxit'
    TT_WECHAT = 'wechat'
    TRANSPORT_TYPES = set([TT_HTTP_API, TT_IRC, TT_TELNET, TT_TWITTER, TT_SMS,
                           TT_USSD, TT_XMPP, TT_MXIT, TT_WECHAT])

    AT_IRC_NICKNAME = 'irc_nickname'
    AT_TELNET_ADDRESS = 'telnet_address'
    AT_TWITTER_HANDLE = 'twitter_handle'
    AT_MSISDN = 'msisdn'
    AT_GTALK_ID = 'gtalk_id'
    AT_JABBER_ID = 'jabber_id'
    AT_MXIT_ID = 'mxit_id'
    AT_WECHAT_ID = 'wechat_id'
    ADDRESS_TYPES = set([
        AT_IRC_NICKNAME, AT_TELNET_ADDRESS, AT_TWITTER_HANDLE, AT_MSISDN,
        AT_GTALK_ID, AT_JABBER_ID, AT_MXIT_ID, AT_WECHAT_ID])

    def process_fields(self, fields):
        fields = super(TransportUserMessage, self).process_fields(fields)
        fields.setdefault('message_id', self.generate_id())
        fields.setdefault('in_reply_to', None)
        fields.setdefault('provider', None)
        fields.setdefault('session_event', None)
        fields.setdefault('content', None)
        fields.setdefault('transport_metadata', {})
        fields.setdefault('group', None)
        fields.setdefault('to_addr_type', None)
        fields.setdefault('from_addr_type', None)
        return fields

    def validate_fields(self):
        super(TransportUserMessage, self).validate_fields()
        # We might get older message versions without the `group` or `provider`
        # fields.
        self.payload.setdefault('group', None)
        self.payload.setdefault('provider', None)
        self.assert_field_present(
            'message_id',
            'to_addr',
            'from_addr',
            'in_reply_to',
            'session_event',
            'content',
            'transport_name',
            'transport_type',
            'transport_metadata',
            'group',
            'provider',
            )
        if self['session_event'] not in self.SESSION_EVENTS:
            raise InvalidMessageField("Invalid session_event %r"
                                      % (self['session_event'],))
        self.assert_field_value('from_addr_type', *self.ADDRESS_TYPES)
        self.assert_field_value('to_addr_type', *self.ADDRESS_TYPES)

    def user(self):
        return self['from_addr']

    def reply(self, content, continue_session=True, **kw):
        """Construct a reply message.

        The reply message will have its `to_addr` field set to the original
        message's `from_addr`. This means that even if the original message is
        directed to the group only (i.e. it has `to_addr` set to `None`), the
        reply will be directed to the sender of the original message.

        :meth:`reply` suitable for constructing both one-to-one messages (such
        as SMS) and directed messages within a group chat (such as
        name-prefixed content in an IRC channel message).

        If `session_event` is provided in the the keyword args,
        `continue_session` will be ignored.

        NOTE: Certain fields are required to come from the message being
              replied to and may not be overridden by this method:

              # If we're not using this addressing, we shouldn't be replying.
              'to_addr', 'from_addr', 'group', 'in_reply_to', 'provider'
              # These three belong together and are supposed to be opaque.
              'transport_name', 'transport_type', 'transport_metadata'

        FIXME: `helper_metadata` should *not* be copied to the reply message.
               We only do it here because a bunch of legacy code relies on it.
        """
        session_event = None if continue_session else self.SESSION_CLOSE

        for field in [
                # If we're not using this addressing, we shouldn't be replying.
                'to_addr', 'from_addr', 'group', 'in_reply_to', 'provider'
                # These three belong together and are supposed to be opaque.
                'transport_name', 'transport_type', 'transport_metadata']:
            if field in kw:
                # Other "bad keyword argument" conditions cause TypeErrors.
                raise TypeError("'%s' may not be overridden." % (field,))

        fields = {
            'helper_metadata': self['helper_metadata'],  # XXX: See above.
            'session_event': session_event,
            'to_addr': self['from_addr'],
            'from_addr': self['to_addr'],
            'group': self['group'],
            'in_reply_to': self['message_id'],
            'provider': self['provider'],
            'transport_name': self['transport_name'],
            'transport_type': self['transport_type'],
            'transport_metadata': self['transport_metadata'],
        }
        fields.update(kw)

        out_msg = TransportUserMessage(content=content, **fields)
        # The reply should go out the same endpoint it came in.
        out_msg.set_routing_endpoint(self.get_routing_endpoint())
        return out_msg

    def reply_group(self, *args, **kw):
        """Construct a group reply message.

        If the `group` field is set to `None`, :meth:`reply_group` is identical
        to :meth:`reply`.

        If the `group` field is not set to `None`, the reply message will have
        its `to_addr` field set to `None`. This means that even if the original
        message is directed to an individual within the group (i.e. its
        `to_addr` is not set to `None`), the reply will be directed to the
        group as a whole.

        :meth:`reply_group` suitable for both one-to-one messages (such as SMS)
        and undirected messages within a group chat (such as IRC channel
        messages).
        """
        out_msg = self.reply(*args, **kw)
        if self['group'] is not None:
            out_msg['to_addr'] = None
        return out_msg

    @classmethod
    def send(cls, to_addr, content, **kw):
        kw.setdefault('from_addr', None)
        kw.setdefault('transport_name', None)
        kw.setdefault('transport_type', None)
        out_msg = cls(
            to_addr=to_addr,
            in_reply_to=None,
            content=content,
            session_event=cls.SESSION_NONE,
            **kw)
        return out_msg

class TransportEvent(TransportMessage):
    """Message about a TransportUserMessage.
    """
    MESSAGE_TYPE = 'event'

    # list of valid delivery statuses
    DELIVERY_STATUSES = frozenset(('pending', 'failed', 'delivered'))

    # map of event_types -> extra fields
    EVENT_TYPES = {
        'ack': {'sent_message_id': lambda v: v is not None},
        'nack': {
            'nack_reason': lambda v: v is not None,
        },
        'delivery_report': {
            'delivery_status': lambda v: v in TransportEvent.DELIVERY_STATUSES,
            },
        }

    def process_fields(self, fields):
        fields = super(TransportEvent, self).process_fields(fields)
        fields.setdefault('event_id', self.generate_id())
        return fields

    def validate_fields(self):
        super(TransportEvent, self).validate_fields()
        self.assert_field_present(
            'user_message_id',
            'event_id',
            'event_type',
            )
        event_type = self.payload['event_type']
        if event_type not in self.EVENT_TYPES:
            raise InvalidMessageField("Unknown event_type %r" % (event_type,))
        for extra_field, check in self.EVENT_TYPES[event_type].items():
            self.assert_field_present(extra_field)
            if not check(self[extra_field]):
                raise InvalidMessageField(extra_field)
<commit_msg>Remove telnet address type.<commit_after># -*- test-case-name: vumi.tests.test_message -*-

import json
from uuid import uuid4
from datetime import datetime

from errors import MissingMessageField, InvalidMessageField

from vumi.utils import to_kwargs

# This is the date format we work with internally
VUMI_DATE_FORMAT = "%Y-%m-%d %H:%M:%S.%f"

def date_time_decoder(json_object):
    for key, value in json_object.items():
        try:
            json_object[key] = datetime.strptime(value,
                    VUMI_DATE_FORMAT)
        except ValueError:
            continue
        except TypeError:
            continue
    return json_object

class JSONMessageEncoder(json.JSONEncoder):
    """A JSON encoder that is able to serialize datetime"""
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime(VUMI_DATE_FORMAT)
        return super(JSONMessageEncoder, self).default(obj)

def from_json(json_string):
    return json.loads(json_string, object_hook=date_time_decoder)

def to_json(obj):
    return json.dumps(obj, cls=JSONMessageEncoder)

class Message(object):
    """
    Start of a somewhat unified message object to be
    used internally in Vumi and while being in transit
    over AMQP

    scary transport format -> Vumi Tansport -> Unified Message -> Vumi Worker

    """

    def __init__(self, _process_fields=True, **kwargs):
        if _process_fields:
            kwargs = self.process_fields(kwargs)
        self.payload = kwargs
        self.validate_fields()

    def process_fields(self, fields):
        return fields

    def validate_fields(self):
        pass

    def assert_field_present(self, *fields):
        for field in fields:
            if field not in self.payload:
                raise MissingMessageField(field)

    def assert_field_value(self, field, *values):
        self.assert_field_present(field)
        if self.payload[field] not in values:
            raise InvalidMessageField(field)

    def to_json(self):
        return to_json(self.payload)

    @classmethod
    def from_json(cls, json_string):
        return cls(_process_fields=False, **to_kwargs(from_json(json_string)))

    def __str__(self):
        return u"<Message payload=\"%s\">" % repr(self.payload)

    def __repr__(self):
        return str(self)

    def __eq__(self, other):
        if isinstance(other, Message):
            return self.payload == other.payload
        return False

    def __contains__(self, key):
        return key in self.payload

    def __getitem__(self, key):
        return self.payload[key]

    def __setitem__(self, key, value):
        self.payload[key] = value

    def get(self, key, default=None):
        return self.payload.get(key, default)

    def items(self):
        return self.payload.items()

    def copy(self):
        return self.from_json(self.to_json())

class TransportMessage(Message):
    """Common base class for messages sent to or from a transport."""

    # sub-classes should set the message type
    MESSAGE_TYPE = None
    MESSAGE_VERSION = '20110921'
    DEFAULT_ENDPOINT_NAME = 'default'

    @staticmethod
    def generate_id():
        """
        Generate a unique message id.

        There are places where we want a message id before we can
        build a complete message. This lets us do that in a consistent
        manner.
        """
        return uuid4().get_hex()

    def process_fields(self, fields):
        fields.setdefault('message_version', self.MESSAGE_VERSION)
        fields.setdefault('message_type', self.MESSAGE_TYPE)
        fields.setdefault('timestamp', datetime.utcnow())
        fields.setdefault('routing_metadata', {})
        fields.setdefault('helper_metadata', {})
        return fields

    def validate_fields(self):
        self.assert_field_value('message_version', self.MESSAGE_VERSION)
        # We might get older event messages without the `helper_metadata`
        # field.
        self.payload.setdefault('helper_metadata', {})
        self.assert_field_present(
            'message_type',
            'timestamp',
            'helper_metadata',
            )
        if self['message_type'] is None:
            raise InvalidMessageField('message_type')

    @property
    def routing_metadata(self):
        return self.payload.setdefault('routing_metadata', {})

    @classmethod
    def check_routing_endpoint(cls, endpoint_name):
        if endpoint_name is None:
            return cls.DEFAULT_ENDPOINT_NAME
        return endpoint_name

    def set_routing_endpoint(self, endpoint_name=None):
        endpoint_name = self.check_routing_endpoint(endpoint_name)
        self.routing_metadata['endpoint_name'] = endpoint_name

    def get_routing_endpoint(self):
        endpoint_name = self.routing_metadata.get('endpoint_name')
        return self.check_routing_endpoint(endpoint_name)

class TransportUserMessage(TransportMessage):
    """Message to or from a user.

    transport_type = sms, ussd, etc
    helper_metadata = for use by dispathers and off-to-the-side
                      components like failure workers (not for use
                      by transports or message workers).
    """

    MESSAGE_TYPE = 'user_message'

    # session event constants
    #
    # SESSION_NONE, SESSION_NEW, SESSION_RESUME, and SESSION_CLOSE
    # may be sent from the transport to a worker. SESSION_NONE indicates
    # there is no relevant session for this message.
    #
    # SESSION_NONE and SESSION_CLOSE may be sent from the worker to
    # the transport. SESSION_NONE indicates any existing session
    # should be continued. SESSION_CLOSE indicates that any existing
    # session should be terminated after sending the message.
    SESSION_NONE, SESSION_NEW, SESSION_RESUME, SESSION_CLOSE = (
        None, 'new', 'resume', 'close')

    # list of valid session events
    SESSION_EVENTS = frozenset([SESSION_NONE, SESSION_NEW, SESSION_RESUME,
                                SESSION_CLOSE])

    # canonical transport types
    TT_HTTP_API = 'http_api'
    TT_IRC = 'irc'
    TT_TELNET = 'telnet'
    TT_TWITTER = 'twitter'
    TT_SMS = 'sms'
    TT_USSD = 'ussd'
    TT_XMPP = 'xmpp'
    TT_MXIT = 'mxit'
    TT_WECHAT = 'wechat'
    TRANSPORT_TYPES = set([TT_HTTP_API, TT_IRC, TT_TELNET, TT_TWITTER, TT_SMS,
                           TT_USSD, TT_XMPP, TT_MXIT, TT_WECHAT])

    AT_IRC_NICKNAME = 'irc_nickname'
    AT_TWITTER_HANDLE = 'twitter_handle'
    AT_MSISDN = 'msisdn'
    AT_GTALK_ID = 'gtalk_id'
    AT_JABBER_ID = 'jabber_id'
    AT_MXIT_ID = 'mxit_id'
    AT_WECHAT_ID = 'wechat_id'
    ADDRESS_TYPES = set([
        AT_IRC_NICKNAME, AT_TWITTER_HANDLE, AT_MSISDN, AT_GTALK_ID,
        AT_JABBER_ID, AT_MXIT_ID, AT_WECHAT_ID])

    def process_fields(self, fields):
        fields = super(TransportUserMessage, self).process_fields(fields)
        fields.setdefault('message_id', self.generate_id())
        fields.setdefault('in_reply_to', None)
        fields.setdefault('provider', None)
        fields.setdefault('session_event', None)
        fields.setdefault('content', None)
        fields.setdefault('transport_metadata', {})
        fields.setdefault('group', None)
        fields.setdefault('to_addr_type', None)
        fields.setdefault('from_addr_type', None)
        return fields

    def validate_fields(self):
        super(TransportUserMessage, self).validate_fields()
        # We might get older message versions without the `group` or `provider`
        # fields.
        self.payload.setdefault('group', None)
        self.payload.setdefault('provider', None)
        self.assert_field_present(
            'message_id',
            'to_addr',
            'from_addr',
            'in_reply_to',
            'session_event',
            'content',
            'transport_name',
            'transport_type',
            'transport_metadata',
            'group',
            'provider',
            )
        if self['session_event'] not in self.SESSION_EVENTS:
            raise InvalidMessageField("Invalid session_event %r"
                                      % (self['session_event'],))
        self.assert_field_value('from_addr_type', *self.ADDRESS_TYPES)
        self.assert_field_value('to_addr_type', *self.ADDRESS_TYPES)

    def user(self):
        return self['from_addr']

    def reply(self, content, continue_session=True, **kw):
        """Construct a reply message.

        The reply message will have its `to_addr` field set to the original
        message's `from_addr`. This means that even if the original message is
        directed to the group only (i.e. it has `to_addr` set to `None`), the
        reply will be directed to the sender of the original message.

        :meth:`reply` suitable for constructing both one-to-one messages (such
        as SMS) and directed messages within a group chat (such as
        name-prefixed content in an IRC channel message).

        If `session_event` is provided in the the keyword args,
        `continue_session` will be ignored.

        NOTE: Certain fields are required to come from the message being
              replied to and may not be overridden by this method:

              # If we're not using this addressing, we shouldn't be replying.
              'to_addr', 'from_addr', 'group', 'in_reply_to', 'provider'
              # These three belong together and are supposed to be opaque.
              'transport_name', 'transport_type', 'transport_metadata'

        FIXME: `helper_metadata` should *not* be copied to the reply message.
               We only do it here because a bunch of legacy code relies on it.
        """
        session_event = None if continue_session else self.SESSION_CLOSE

        for field in [
                # If we're not using this addressing, we shouldn't be replying.
                'to_addr', 'from_addr', 'group', 'in_reply_to', 'provider'
                # These three belong together and are supposed to be opaque.
                'transport_name', 'transport_type', 'transport_metadata']:
            if field in kw:
                # Other "bad keyword argument" conditions cause TypeErrors.
                raise TypeError("'%s' may not be overridden." % (field,))

        fields = {
            'helper_metadata': self['helper_metadata'],  # XXX: See above.
            'session_event': session_event,
            'to_addr': self['from_addr'],
            'from_addr': self['to_addr'],
            'group': self['group'],
            'in_reply_to': self['message_id'],
            'provider': self['provider'],
            'transport_name': self['transport_name'],
            'transport_type': self['transport_type'],
            'transport_metadata': self['transport_metadata'],
        }
        fields.update(kw)

        out_msg = TransportUserMessage(content=content, **fields)
        # The reply should go out the same endpoint it came in.
        out_msg.set_routing_endpoint(self.get_routing_endpoint())
        return out_msg

    def reply_group(self, *args, **kw):
        """Construct a group reply message.

        If the `group` field is set to `None`, :meth:`reply_group` is identical
        to :meth:`reply`.

        If the `group` field is not set to `None`, the reply message will have
        its `to_addr` field set to `None`. This means that even if the original
        message is directed to an individual within the group (i.e. its
        `to_addr` is not set to `None`), the reply will be directed to the
        group as a whole.

        :meth:`reply_group` suitable for both one-to-one messages (such as SMS)
        and undirected messages within a group chat (such as IRC channel
        messages).
        """
        out_msg = self.reply(*args, **kw)
        if self['group'] is not None:
            out_msg['to_addr'] = None
        return out_msg

    @classmethod
    def send(cls, to_addr, content, **kw):
        kw.setdefault('from_addr', None)
        kw.setdefault('transport_name', None)
        kw.setdefault('transport_type', None)
        out_msg = cls(
            to_addr=to_addr,
            in_reply_to=None,
            content=content,
            session_event=cls.SESSION_NONE,
            **kw)
        return out_msg

class TransportEvent(TransportMessage):
    """Message about a TransportUserMessage.
    """
    MESSAGE_TYPE = 'event'

    # list of valid delivery statuses
    DELIVERY_STATUSES = frozenset(('pending', 'failed', 'delivered'))

    # map of event_types -> extra fields
    EVENT_TYPES = {
        'ack': {'sent_message_id': lambda v: v is not None},
        'nack': {
            'nack_reason': lambda v: v is not None,
        },
        'delivery_report': {
            'delivery_status': lambda v: v in TransportEvent.DELIVERY_STATUSES,
            },
        }

    def process_fields(self, fields):
        fields = super(TransportEvent, self).process_fields(fields)
        fields.setdefault('event_id', self.generate_id())
        return fields

    def validate_fields(self):
        super(TransportEvent, self).validate_fields()
        self.assert_field_present(
            'user_message_id',
            'event_id',
            'event_type',
            )
        event_type = self.payload['event_type']
        if event_type not in self.EVENT_TYPES:
            raise InvalidMessageField("Unknown event_type %r" % (event_type,))
        for extra_field, check in self.EVENT_TYPES[event_type].items():
            self.assert_field_present(extra_field)
            if not check(self[extra_field]):
                raise InvalidMessageField(extra_field)
<|endoftext|>