<commit_before>import functools
import inspect

from .particle_class import Particle
from ..utils import AtomicError

# TODO: make sure particle_input works with classes and class methods

def particle_input(wrapped_function=None, **kwargs):
    r"""A decorator to take arguments and keywords related to a
    particle, and pass through the Particle class to the callable instead.

    Notes
    -----

    This version of particle_input only works with functions, but in the future
    should be extended to work with classes and methods.

    """

    def decorator(wrapped_function):
        wrapped_signature = inspect.signature(wrapped_function)

        @functools.wraps(wrapped_function)
        def wrapper(*args, **kwargs):

            # The following couple of statements will likely need to be
            # modified in order to work with methods.

            bound_args = wrapped_signature.bind(*args, **kwargs)
            arguments = bound_args.arguments
            argnames = bound_args.arguments.keys()

            annotations = wrapped_function.__annotations__

            args_to_become_particles = [
                argname for argname in annotations.keys()
                if annotations[argname] is Particle
            ]

            if not args_to_become_particles:
                raise AtomicError(
                    f"None of the arguments or keywords to {func.__name__} "
                    f"have been annotated with the Particle class as required "
                    f"by the particle_input decorator. For example,\n\n"
                    f"    def func(arg, particle: Particle, kwarg=None):\n\n"
                    f"will convert the particle argument to ")

            # If the number of arguments and keywords annotated with
            # Particle is exactly one, then the Z and mass_numb keywords
            # can be used without potential for ambiguity.

            if len(args_to_become_particles) == 1:
                Z = arguments.get('Z', None)
                mass_numb = arguments.get('mass_numb', None)
            else:
                if 'Z' in argnames or 'mass_numb' in argnames:
                    raise AtomicError(
                        f"The arguments Z and mass_numb in "
                        f"{wrapped_function.__name__} are not allowed when "
                        f"more than one argument or keyword is annotated "
                        f"with Particle in order to avoid ambiguity.")
                else:
                    Z = None
                    mass_numb = None

            if Z is not None and not isinstance(Z, int):
                raise TypeError(
                    f"The argument Z = {repr(Z)} in "
                    f"{wrapped_function.__name__} is not an integer.")

            if mass_numb is not None and not isinstance(mass_numb, int):
                raise TypeError(
                    f"The argument mass_numb in = {repr(mass_numb)} "
                    f"{wrapped_function.__name__} is not an integer.")

            # Go through the argument names and check whether or not they are
            # annotated with Particle.  If they aren't, include the name and
            # value of the argument as an item in the new keyword arguments
            # dictionary unchanged.  If they are annotated with Particle, then
            # either convert the representation of a Particle to a Particle if
            # it is not already a Particle.

            new_kwargs = {}

            for argname in argnames:
                argval = arguments[argname]

                should_be_particle = argname in args_to_become_particles
                not_already_particle = not isinstance(argval, Particle)

                if should_be_particle and not_already_particle:
                    new_kwargs[argname] = \
                        Particle(argval, Z=Z, mass_numb=mass_numb)
                else:
                    new_kwargs[argname] = argval

            return wrapped_function(**new_kwargs)

        return wrapper

    # The following code helps allow the decorator to be used either with
    # or without arguments.  In particular, this helps allow us to invoke
    # the decorator either as `@particle_input` or as `@particle_input()`,
    # where the latter call allows the decorator to have keyword arguments.

    if wrapped_function is not None:
        return decorator(wrapped_function)
    else:
        return decorator
<commit_msg>Update @particle_input<commit_after>import functools
import inspect

from .particle_class import Particle
from ..utils import (AtomicError,
                     InvalidParticleError,
                     InvalidElementError,
                     InvalidIonError,
                     InvalidIsotopeError,
                     ChargeError)

from typing import Callable, Union, Any, Set, List, Tuple

# TODO: make sure particle_input works with classes and class methods

def particle_input(wrapped_function: Callable = None,
                   must_be: Union[str, Set, List, Tuple] = set(),
                   cannot_be: Union[str, Set, List, Tuple] = set(),
                   any: bool = False,
                   **kwargs) -> Any:
    r"""A decorator to take arguments and keywords related to a
    particle, and pass through the Particle class to the callable instead.

    Parameters
    ----------

    wrapped_function : callable
        The function to be decorated.

    must_be : str, set, list, or tuple; optional
        The categories that the Particle must be in; otherwise the

    Notes
    -----

    This version of particle_input only works with functions, but in the future
    should be extended to work with classes and methods.

    """

    def _errmsg(argname: str, argval: str, Z: int = None,
                mass_numb: int = None, funcname: str = None) -> str:
        r"""Returns a string with an appropriate error message for an
        InvalidParticleError."""

        errmsg = f"In {funcname}, {argname} = {repr(argval)} "
        if mass_numb is not None or Z is not None:
            errmsg += "with "
        if mass_numb is not None:
            errmsg += f"mass_numb = {repr(mass_numb)} "
        if mass_numb is not None and Z is not None:
            errmsg += "and "
        if Z is not None:
            errmsg += f"integer charge Z = {repr(Z)} "
        errmsg += "does not correspond to a valid particle."
        return errmsg

    def decorator(wrapped_function: Callable):
        wrapped_signature = inspect.signature(wrapped_function)

        @functools.wraps(wrapped_function)
        def wrapper(*args, **kwargs):

            # The following couple of statements will likely need to be
            # modified in order to work with methods.

            bound_args = wrapped_signature.bind(*args, **kwargs)
            arguments = bound_args.arguments
            argnames = bound_args.arguments.keys()
            funcname = wrapped_function.__name__
            annotations = wrapped_function.__annotations__

            args_to_become_particles = [
                argname for argname in annotations.keys()
                if annotations[argname] is Particle
            ]

            if not args_to_become_particles:
                raise AtomicError(
                    f"None of the arguments or keywords to {funcname} "
                    f"have been annotated with Particle as required "
                    f"by the particle_input decorator.")
            elif len(args_to_become_particles) > 1:
                if 'Z' in argnames or 'mass_numb' in argnames:
                    raise AtomicError(
                        f"The arguments Z and mass_numb in {funcname} are not "
                        f"allowed when more than one argument or keyword is "
                        f"annotated with Particle in functions decorated "
                        f"with @particle_input.")

            # If the number of arguments and keywords annotated with
            # Particle is exactly one, then the Z and mass_numb keywords
            # can be used without potential for ambiguity.

            Z = arguments.get('Z', None)
            mass_numb = arguments.get('mass_numb', None)

            if Z is not None and not isinstance(Z, int):
                raise TypeError(
                    f"The argument Z = {repr(Z)} in {funcname} is not an "
                    f"integer.")

            if mass_numb is not None and not isinstance(mass_numb, int):
                raise TypeError(
                    f"The argument mass_numb = {repr(mass_numb)} in "
                    f"{funcname} is not an integer.")

            # Go through the argument names and check whether or not they are
            # annotated with Particle.  If they aren't, include the name and
            # value of the argument as an item in the new keyword arguments
            # dictionary unchanged.  If they are annotated with Particle, then
            # either convert the representation of a Particle to a Particle if
            # it is not already a Particle.

            new_kwargs = {}

            for argname in argnames:
                argval = arguments[argname]

                should_be_particle = argname in args_to_become_particles
                not_already_particle = not isinstance(argval, Particle)

                if should_be_particle and not_already_particle:

                    if not isinstance(argval, (int, str)):
                        raise TypeError

                    try:
                        particle = Particle(argval, Z=Z, mass_numb=mass_numb)
                    except InvalidParticleError as e:
                        raise InvalidParticleError(
                            _errmsg(argname, argval, Z, mass_numb, funcname)
                        ) from e

                    if argname == 'element' and not particle.element:
                        raise InvalidElementError
                    if argname == 'isotope' and not particle.isotope:
                        raise InvalidIsotopeError
                    if argname == 'ion' and not particle.ion:
                        raise InvalidIonError
                    if 'charged' in must_be and not particle._integer_charge:
                        raise ChargeError
                    if not particle.is_category(must_be, exclude=cannot_be):
                        raise AtomicError

                    new_kwargs[argname] = particle
                else:
                    new_kwargs[argname] = argval

            return wrapped_function(**new_kwargs)

        return wrapper

    # The following code helps allow the decorator to be used either with
    # or without arguments.  In particular, this helps allow us to invoke
    # the decorator either as `@particle_input` or as `@particle_input()`,
    # where the latter call allows the decorator to have keyword arguments.

    if wrapped_function is not None:
        return decorator(wrapped_function)
    else:
        return decorator
<|endoftext|>


<commit_before>#encoding=utf-8
import unittest
from model import piece

class PieceCase(unittest.TestCase):
    def setUp(self):
        pass
    def tearDown(self):
        pass

    def testNoAuthor(self):
        # 普通句子
        self.assertEqual(piece.parse_content(u"每个人的心里，有多么长的一个清单，这些清单里写着多少美好的事，可是，它们总是被推迟，被搁置，在时间的阁楼上腐烂。为什么勇气的问题总是被误以为是时间的问题，而那些沉重的、抑郁的、不得已的，总是被叫做生活本身。"),{
            "content":u"每个人的心里，有多么长的一个清单，这些清单里写着多少美好的事，可是，它们总是被推迟，被搁置，在时间的阁楼上腐烂。为什么勇气的问题总是被误以为是时间的问题，而那些沉重的、抑郁的、不得已的，总是被叫做生活本身。",
            "author":None,
            "work":None
        })

    def testPureEnglishWithOneWord(self):
        self.assertEqual(piece.parse_content(u"Mac-talk."),{
            "content":u"Mac-talk.",
            "author":None,
            "work":None
        })

    def testTooManySpliter(self):
        # 过多链接符忽略
        self.assertEqual(piece.parse_content("https://stacksocial.com/sales/the-mac-freebie-bundle-2-0"),{
            "content":"https://stacksocial.com/sales/the-mac-freebie-bundle-2-0",
            "author":None,
            "work":None
        })

    def testByIssue(self):

        self.assertEqual(piece.parse_content(u"It's better to fail than to pass by pure luck."),{
            "content":u"It's better to fail than to pass by pure luck.",
            "author":None,
            "work":None
        })

    def testAuthorSentance(self):
        self.assertEqual(piece.parse_content(u"我会唱歌我知道\n对于评审给我这样的肯定\n我也给予肯定\n----王菲，金曲獎十五屆最佳女演唱人得獎感言"),{
            "content":u"我会唱歌我知道\n对于评审给我这样的肯定\n我也给予肯定\n----王菲，金曲獎十五屆最佳女演唱人得獎感言",
            "author":None,
            "work":None
        })

    def testParseContent(self):
        """ testParseContent """
        self.assertEqual(piece.parse_content(u"愚蠢的人饱受其愚蠢所带来的疲累之苦。       by 叔本华"),{
            "content":u"愚蠢的人饱受其愚蠢所带来的疲累之苦。",
            "author":u"叔本华",
            "work":None
        })

        # 带言说者
        self.assertEqual(piece.parse_content(u"我们终将浑然一体，就像水溶于水。 \n——柴静"),{
            "content":u"我们终将浑然一体，就像水溶于水。",
            "author":u'柴静',
            "work":None
        })
        # 不规范破折 ------
        self.assertEqual(piece.parse_content(u"黄金真是一个奇妙的东西！谁拥有了它，谁就成为他想要的一切东西的主人，有了黄金，甚至可以使灵魂升入天堂。-------哥伦布"),{
            "content":u"黄金真是一个奇妙的东西！谁拥有了它，谁就成为他想要的一切东西的主人，有了黄金，甚至可以使灵魂升入天堂。",
            "author":u"哥伦布",
            "work":None
        })
        # 人名
        self.assertEqual(piece.parse_content(u"当下最流行三大病症：拖延症、强迫症、选择困难症，直白点说就是懒、贱、穷。----@王黎wario"),{
            "content":u"当下最流行三大病症：拖延症、强迫症、选择困难症，直白点说就是懒、贱、穷。",
            "author":u"@王黎wario",
            "work":None
        })
        # 书名
        self.assertEqual(piece.parse_content(u"我一直坚持的一个信念是，改变不了大环境，就改变小环境，做自己力所能及的事情。你不能决定太阳几点升起，但可以决定自己几点起床。——《自由在高处》"),{
            "content":u"我一直坚持的一个信念是，改变不了大环境，就改变小环境，做自己力所能及的事情。你不能决定太阳几点升起，但可以决定自己几点起床。",
            "author":None,
            "work":u"自由在高处"
        })
        # 人名及书名
        self.assertEqual(piece.parse_content(u"有时候，你很想念一个人，但你不会打电话给他。打电话给他，不知道说甚么好，还是不打比较好。想念一个人，不一定要听到他的声音。听到了他的声音也许就是另一回事。 想像中的一切往往比现实稍微美好一点。想念中的那个人也比现实稍微温暖一点。—— 约翰·肖尔斯《许愿树》"),{
            "content":u"有时候，你很想念一个人，但你不会打电话给他。打电话给他，不知道说甚么好，还是不打比较好。想念一个人，不一定要听到他的声音。听到了他的声音也许就是另一回事。 想像中的一切往往比现实稍微美好一点。想念中的那个人也比现实稍微温暖一点。",
            "author":u"约翰·肖尔斯",
            "work":u"许愿树"
        })
        # 英文人名书名
        self.assertEqual(piece.parse_content(u"To choose doubt as a philosophy of life is akin to choosing immobility as a means of transportation. —— Yann Martel 《Life of Pi》"),{
            "content":u"To choose doubt as a philosophy of life is akin to choosing immobility as a means of transportation.",
            "author":u"Yann Martel",
            "work":u"Life of Pi"
        })
        # 英文人名
        self.assertEqual(piece.parse_content(u"qweqweq qwe q we qw eqw ew —— qweqwe"),{
            "content":u"qweqweq qwe q we qw eqw ew",
            "author":u"qweqwe",
            "work":None
        })
        # 过长中文认为非人名
        self.assertEqual(piece.parse_content(u"曰：我和我的小伙伴们都惊呆了？\n\n　　\n\n　　问：为什么惊呆了？\n\n　　--这个该怎么回答"),{
            "content":u"曰：我和我的小伙伴们都惊呆了？\n\n　　\n\n　　问：为什么惊呆了？\n\n　　--这个该怎么回答",
            "author":None,
            "work":None
        })<commit_msg>add case<commit_after>#encoding=utf-8
import unittest
from model import piece

class PieceCase(unittest.TestCase):
    def setUp(self):
        pass
    def tearDown(self):
        pass

    def testNoAuthor(self):
        # 普通句子
        self.assertEqual(piece.parse_content(u"每个人的心里，有多么长的一个清单，这些清单里写着多少美好的事，可是，它们总是被推迟，被搁置，在时间的阁楼上腐烂。为什么勇气的问题总是被误以为是时间的问题，而那些沉重的、抑郁的、不得已的，总是被叫做生活本身。"),{
            "content":u"每个人的心里，有多么长的一个清单，这些清单里写着多少美好的事，可是，它们总是被推迟，被搁置，在时间的阁楼上腐烂。为什么勇气的问题总是被误以为是时间的问题，而那些沉重的、抑郁的、不得已的，总是被叫做生活本身。",
            "author":None,
            "work":None
        })

    def testPureEnglishWithOneWord(self):
        self.assertEqual(piece.parse_content(u"Mac-talk."),{
            "content":u"Mac-talk.",
            "author":None,
            "work":None
        })

    def testTooManySpliter(self):
        # 过多链接符忽略
        self.assertEqual(piece.parse_content("https://stacksocial.com/sales/the-mac-freebie-bundle-2-0"),{
            "content":"https://stacksocial.com/sales/the-mac-freebie-bundle-2-0",
            "author":None,
            "work":None
        })

    def testByIssue(self):

        self.assertEqual(piece.parse_content(u"It's better to fail than to pass by pure luck."),{
            "content":u"It's better to fail than to pass by pure luck.",
            "author":None,
            "work":None
        })

    def testAuthorSentance(self):
        self.assertEqual(piece.parse_content(u"我会唱歌我知道\n对于评审给我这样的肯定\n我也给予肯定\n----王菲，金曲獎十五屆最佳女演唱人得獎感言"),{
            "content":u"我会唱歌我知道\n对于评审给我这样的肯定\n我也给予肯定\n----王菲，金曲獎十五屆最佳女演唱人得獎感言",
            "author":None,
            "work":None
        })

    def testParseContent(self):
        """ testParseContent """
        self.assertEqual(piece.parse_content(u"愚蠢的人饱受其愚蠢所带来的疲累之苦。       by 叔本华"),{
            "content":u"愚蠢的人饱受其愚蠢所带来的疲累之苦。",
            "author":u"叔本华",
            "work":None
        })

        # 带言说者
        self.assertEqual(piece.parse_content(u"我们终将浑然一体，就像水溶于水。 \n——柴静"),{
            "content":u"我们终将浑然一体，就像水溶于水。",
            "author":u'柴静',
            "work":None
        })
        # 不规范破折 ------
        self.assertEqual(piece.parse_content(u"黄金真是一个奇妙的东西！谁拥有了它，谁就成为他想要的一切东西的主人，有了黄金，甚至可以使灵魂升入天堂。-------哥伦布"),{
            "content":u"黄金真是一个奇妙的东西！谁拥有了它，谁就成为他想要的一切东西的主人，有了黄金，甚至可以使灵魂升入天堂。",
            "author":u"哥伦布",
            "work":None
        })
        # 人名
        self.assertEqual(piece.parse_content(u"当下最流行三大病症：拖延症、强迫症、选择困难症，直白点说就是懒、贱、穷。----@王黎wario"),{
            "content":u"当下最流行三大病症：拖延症、强迫症、选择困难症，直白点说就是懒、贱、穷。",
            "author":u"@王黎wario",
            "work":None
        })
        # 书名
        self.assertEqual(piece.parse_content(u"我一直坚持的一个信念是，改变不了大环境，就改变小环境，做自己力所能及的事情。你不能决定太阳几点升起，但可以决定自己几点起床。——《自由在高处》"),{
            "content":u"我一直坚持的一个信念是，改变不了大环境，就改变小环境，做自己力所能及的事情。你不能决定太阳几点升起，但可以决定自己几点起床。",
            "author":None,
            "work":u"自由在高处"
        })
        # 人名及书名
        self.assertEqual(piece.parse_content(u"有时候，你很想念一个人，但你不会打电话给他。打电话给他，不知道说甚么好，还是不打比较好。想念一个人，不一定要听到他的声音。听到了他的声音也许就是另一回事。 想像中的一切往往比现实稍微美好一点。想念中的那个人也比现实稍微温暖一点。—— 约翰·肖尔斯《许愿树》"),{
            "content":u"有时候，你很想念一个人，但你不会打电话给他。打电话给他，不知道说甚么好，还是不打比较好。想念一个人，不一定要听到他的声音。听到了他的声音也许就是另一回事。 想像中的一切往往比现实稍微美好一点。想念中的那个人也比现实稍微温暖一点。",
            "author":u"约翰·肖尔斯",
            "work":u"许愿树"
        })
        self.assertEqual(piece.parse_content(u"永远不要忘记，直至上帝向人揭示出未来之日，人类全部智慧就包含在两个词中：等待和希望。 —《基督山伯爵》"),{
            "content":u"永远不要忘记，直至上帝向人揭示出未来之日，人类全部智慧就包含在两个词中：等待和希望。",
            "author": None,
            "work": u"基督山伯爵"
        })
        # 英文人名书名
        self.assertEqual(piece.parse_content(u"To choose doubt as a philosophy of life is akin to choosing immobility as a means of transportation. —— Yann Martel 《Life of Pi》"),{
            "content":u"To choose doubt as a philosophy of life is akin to choosing immobility as a means of transportation.",
            "author":u"Yann Martel",
            "work":u"Life of Pi"
        })
        # 英文人名
        self.assertEqual(piece.parse_content(u"qweqweq qwe q we qw eqw ew —— qweqwe"),{
            "content":u"qweqweq qwe q we qw eqw ew",
            "author":u"qweqwe",
            "work":None
        })
        # 过长中文认为非人名
        self.assertEqual(piece.parse_content(u"曰：我和我的小伙伴们都惊呆了？\n\n　　\n\n　　问：为什么惊呆了？\n\n　　--这个该怎么回答"),{
            "content":u"曰：我和我的小伙伴们都惊呆了？\n\n　　\n\n　　问：为什么惊呆了？\n\n　　--这个该怎么回答",
            "author":None,
            "work":None
        })<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
    prestashop

    :copyright: (c) 2015 by Openlabs Technologies & Consulting (P) Limited
    :license: GPLv3, see LICENSE for more details.
"""
from datetime import datetime

import pytz
import requests
import pystashop
from mockstashop import MockstaShopWebservice
from trytond.model import ModelView, fields
from trytond.transaction import Transaction
from trytond.pool import Pool, PoolMeta
from trytond.wizard import Wizard, StateView, Button
from trytond.pyson import Eval

__metaclass__ = PoolMeta
__all__ = [
    'Channel', 'PrestashopImportOrdersWizardView',
    'PrestashopImportOrdersWizard', 'PrestashopExportOrdersWizardView',
    'PrestashopExportOrdersWizard',
    'PrestashopConnectionWizardView', 'PrestashopConnectionWizard',
]
TIMEZONES = [(None, '')] + [(x, x) for x in pytz.common_timezones]

PRESTASHOP_STATES = {
    'required': Eval('source') == 'prestashop',
    'invisible': ~(Eval('source') == 'prestashop')
}
INVISIBLE_IF_NOT_PRESTASHOP = {
    'invisible': ~(Eval('source') == 'prestashop')
}

class Channel:
    """
    Sale Channel model
    """
    __name__ = 'sale.channel'

    #: The URL of prestashop site
    prestashop_url = fields.Char(
        'Prestashop URL', states=PRESTASHOP_STATES, depends=['source']
    )

    #: The webservices key for access to site
    prestashop_key = fields.Char(
        'Prestashop Key', states=PRESTASHOP_STATES, depends=['source']
    )

    #: Last time at which the orders were imported from prestashop
    last_prestashop_order_import_time = fields.DateTime(
        'Last Prestashop Order Import Time', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )

    #: Last time at which the orders were exported to prestashop
    last_prestashop_order_export_time = fields.DateTime(
        'Last Prestashop order export time', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )
    prestashop_shipping_product = fields.Many2One(
        'product.product', 'Shipping Product', states=PRESTASHOP_STATES,
        domain=[
            ('type', '=', 'service'),
            ('template.type', '=', 'service'),
        ], depends=['source']
    )

    #: The timezone set on prestashop site
    #: The orders imported will bear time in this timezone and that will need
    #: to be converted to UTC
    #: Also in order to determine what orders are to be imported, we need
    #: to convert UTC to this timezone to ensure correct time interval
    prestashop_timezone = fields.Selection(
        TIMEZONES, 'Timezone', translate=False, states=PRESTASHOP_STATES,
        depends=['source']
    )

    #: Allowed languages to be synced for the site
    prestashop_languages = fields.One2Many(
        'prestashop.site.lang', 'channel', 'Languages',
        states=INVISIBLE_IF_NOT_PRESTASHOP, depends=['source']
    )

    #: The mapping between prestashop order states and tryton sale states
    prestashop_order_states = fields.One2Many(
        'prestashop.site.order_state', 'channel', 'Order States',
        states=INVISIBLE_IF_NOT_PRESTASHOP, depends=['source']
    )

    #: Set this to True to handle invoicing in tryton and get invoice info
    #: TODO: A provision to be implemented in future versions
    #: Also handle multiple payment methods where each will have different
    #: journal and account setups
    prestashop_handle_invoice = fields.Boolean(
        'Handle Invoicing ?', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )

    @classmethod
    def get_source(cls):
        """
        Get the source
        """
        res = super(Channel, cls).get_source()
        res.append(('prestashop', 'Prestashop'))
        return res

    def validate_prestashop_channel(self):
        """
        Check if current channel belongs to prestashop
        """
        if self.source != 'prestashop':
            self.raise_user_error("invalid_prestashop_channel")

    @classmethod
    def __setup__(cls):
        super(Channel, cls).__setup__()
        cls._error_messages.update({
            'prestashop_settings_missing':
                'Prestashop webservice settings are incomplete.',
            'invalid_prestashop_channel':
                'Current channel does not belongs to prestashop',
            'multiple_channels':
                'Test connection can be done for only one channel at a time.',
            'wrong_url_n_key': 'Connection Failed! Please check URL and Key',
            'wrong_url': 'Connection Failed! The URL provided is wrong',
            'languages_not_imported':
                'Import the languages before importing order states',
            'order_states_not_imported':
                'Import the order states before importing/exporting orders'
        })
        cls._buttons.update({
            'test_prestashop_connection': {},
            'import_prestashop_languages': {},
            'import_prestashop_order_states': {},
            'import_prestashop_orders': {},
            'export_prestashop_orders': {},
        })

    def get_prestashop_client(self):
        """
        Returns an authenticated instance of the Prestashop client

        :return: Prestashop client object
        """
        if not all([self.prestashop_url, self.prestashop_key]):
            self.raise_user_error('prestashop_settings_missing')

        if Transaction().context.get('ps_test'):
            return MockstaShopWebservice('Some URL', 'A Key')

        return pystashop.PrestaShopWebservice(
            self.prestashop_url, self.prestashop_key
        )

    @classmethod
    @ModelView.button
    def import_prestashop_languages(cls, channels):
        """Import Languages from remote and try to link them to tryton
        languages

        :param sites: List of sites for which the languages are to be imported
        :returns: List of languages created
        """
        SiteLanguage = Pool().get('prestashop.site.lang')

        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()

        # Set this site in context
        with Transaction().set_context(current_channel=channel.id):

            client = channel.get_prestashop_client()
            languages = client.languages.get_list(display='full')

            new_records = []
            for lang in languages:
                # If the language already exists in `Languages`, skip and do
                # not create it again
                if SiteLanguage.search_using_ps_id(lang.id.pyval):
                    continue
                new_records.append(
                    SiteLanguage.create_using_ps_data(lang)
                )

        return new_records

    @classmethod
    @ModelView.button
    def import_prestashop_order_states(cls, channels):
        """Import Order States from remote and try to link them to tryton
        order states

        :returns: List of order states created
        """
        SiteOrderState = Pool().get('prestashop.site.order_state')

        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()

        # Set this channel to context
        with Transaction().set_context(current_channel=channel.id):

            # If channel languages don't exist, then raise an error
            if not channel.prestashop_languages:
                cls.raise_user_error('languages_not_imported')

            client = channel.get_prestashop_client()
            order_states = client.order_states.get_list(display='full')

            new_records = []
            for state in order_states:
                # If this order state already exists for this channel, skip and
                # do not create it again
                if SiteOrderState.search_using_ps_id(state.id.pyval):
                    continue

                new_records.append(
                    SiteOrderState.create_using_ps_data(state)
                )

        return new_records

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_connection')
    def test_prestashop_connection(cls, channels):
        """Test Prestashop connection and display appropriate message to user
        """
        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()
        try:
            client = channel.get_prestashop_client()

            # Try getting the list of shops
            # If it fails with prestashop error, then raise error
            client.shops.get_list()
        except pystashop.PrestaShopWebserviceException:
            cls.raise_user_error('wrong_url_n_key')
        except (
            requests.exceptions.MissingSchema,
            requests.exceptions.ConnectionError
        ):
            cls.raise_user_error('wrong_url')

    @classmethod
    def import_prestashop_orders_using_cron(cls, channels):
        """Import orders from prestashop

        :param sites: The list of sites from which the orders are to be
                      imported

        ..note:: This method is usually called by the cron
        """
        channels = cls.search([
            ('source', '=', 'prestashop')
        ])
        for channel in channels:
            channel.import_prestashop_orders()

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_import_orders')
    def import_prestashop_orders_button(cls, sites=None):
        """Dummy button to fire up the wizard for import of orders

        :param sites: The list of sites from which the orders are to be
                      imported
        """
        pass

    def import_prestashop_orders(self):
        """Import orders for the current site from prestashop
        Import only those orers which are updated after the
        `last order import time` as set in the prestashop configuration

        :returns: The list of active records of sales imported
        """
        Sale = Pool().get('sale.sale')
        self.validate_prestashop_channel()

        if not self.prestashop_order_states:
            self.raise_user_error('order_states_not_imported')

        # Localize to the site timezone
        utc_time_now = datetime.utcnow()
        site_tz = pytz.timezone(self.prestashop_timezone)
        time_now = site_tz.normalize(pytz.utc.localize(utc_time_now))
        client = self.get_prestashop_client()

        with Transaction().set_context(current_channel=self.id):
            if self.last_prestashop_order_import_time:
                # In tryton all time stored is in UTC
                # Convert the last import time to timezone of the site
                last_prestashop_order_import_time = site_tz.normalize(
                    pytz.utc.localize(self.last_prestashop_order_import_time)
                )
                orders_to_import = client.orders.get_list(
                    filters={
                        'date_upd': '{0},{1}'.format(
                            last_prestashop_order_import_time.strftime(
                                '%Y-%m-%d %H:%M:%S'
                            ),
                            time_now.strftime('%Y-%m-%d %H:%M:%S')
                        )
                    }, date=1, display='full'
                )
            else:
                # FIXME: This wont scale if there are thousands of orders
                orders_to_import = client.orders.get_list(display='full')

            self.write([self], {
                'last_prestashop_order_import_time': utc_time_now
            })
            sales_imported = []
            for order in orders_to_import:
                sales_imported.append(Sale.find_or_create_using_ps_data(order))

        return sales_imported

    @classmethod
    def export_orders_to_prestashop_using_cron(cls):
        """
        Export order status to prestashop using cron
        """
        channels = cls.search([
            ('source', '=', 'prestashop')
        ])
        for channel in channels:
            channel.export_orders_to_prestashop()

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_export_orders')
    def export_orders_button(cls, channels):
        """
        Dummy button to fire up the wizard for export of orders
        """
        pass

    def export_orders_to_prestashop(self):
        """
        Export order status to prestashop current site
        Export only those orders which are modified after the
        `last order export time` as set in the prestashop configuration.

        :returns: The list of active records of sales exported
        """
        Sale = Pool().get('sale.sale')
        Move = Pool().get('stock.move')

        if not self.order_states:
            self.raise_user_error('order_states_not_imported')

        time_now = datetime.utcnow()

        self.validate_prestashop_channel()

        with Transaction().set_context(current_channel=self.id):
            if self.last_prestashop_order_export_time:
                # Sale might not get updated for state changes in the related
                # shipments.
                # So first get the moves for outgoing shipments which are \
                # executed after last import time.
                moves = Move.search([
                    (
                        'write_date', '>=',
                        self.last_prestashop_order_export_time
                    ),
                    ('sale.channel', '=', self.id),
                    ('shipment', 'like', 'stock.shipment.out%')
                ])
                sales_to_export = Sale.search(['OR', [
                    ('write_date', '>=', self.last_order_export_time),
                    ('channel', '=', self.id),
                ], [
                    ('id', 'in', map(int, [m.sale for m in moves]))
                ]])
            else:
                sales_to_export = Sale.search([('channel', '=', self.id)])

            self.write([self], {
                'last_prestashop_order_export_time': time_now
            })

            for sale in sales_to_export:
                sale.export_status_to_ps()

        return sales_to_export

class PrestashopConnectionWizardView(ModelView):
    'Prestashop Connection Wizard View'
    __name__ = 'prestashop.connection.wizard.view'

class PrestashopConnectionWizard(Wizard):
    'Prestashop Connection Wizard'
    __name__ = 'prestashop.connection.wizard'

    start = StateView(
        'prestashop.connection.wizard.view',
        'prestashop.prestashop_connection_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, data):
        """Test the connection and show the user appropriate message

        :param data: Wizard data
        """
        return {}

class PrestashopImportOrdersWizardView(ModelView):
    'Prestashop Import Wizard View'
    __name__ = 'prestashop.import_orders.wizard.view'

    orders_imported = fields.Integer('Orders Imported', readonly=True)

class PrestashopImportOrdersWizard(Wizard):
    'Prestashop Import Wizard'
    __name__ = 'prestashop.import_orders.wizard'

    start = StateView(
        'prestashop.import_orders.wizard.view',
        'prestashop.prestashop_import_orders_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, fields):
        """
        Import the orders and display a confirmation message to the user

        :param fields: Wizard fields
        """
        SaleChannel = Pool().get('sale.channel')

        channel = SaleChannel(Transaction().context['active_id'])

        channel.validate_prestashop_channel()

        return {
            'orders_imported': len(channel.import_prestashop_orders())
        }

class PrestashopExportOrdersWizardView(ModelView):
    'Prestashop Export Orders Wizard View'
    __name__ = 'prestashop.export_orders.wizard.view'

    orders_exported = fields.Integer('Orders Exported', readonly=True)

class PrestashopExportOrdersWizard(Wizard):
    'Prestashop Export Orders Wizard'
    __name__ = 'prestashop.export_orders.wizard'

    start = StateView(
        'prestashop.export_orders.wizard.view',
        'prestashop.prestashop_export_orders_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, fields):
        """
        Export the orders and display a confirmation message to the user

        :param fields: Wizard fields
        """
        SaleChannel = Pool().get('sale.channel')

        channel = SaleChannel(Transaction().context['active_id'])

        channel.validate_prestashop_channel()

        default = {
            'orders_exported': len(channel.export_orders_to_prestashop())
        }

        return default
<commit_msg>Shipping product must be salable #7696<commit_after># -*- coding: utf-8 -*-
"""
    prestashop

    :copyright: (c) 2015 by Openlabs Technologies & Consulting (P) Limited
    :license: GPLv3, see LICENSE for more details.
"""
from datetime import datetime

import pytz
import requests
import pystashop
from mockstashop import MockstaShopWebservice
from trytond.model import ModelView, fields
from trytond.transaction import Transaction
from trytond.pool import Pool, PoolMeta
from trytond.wizard import Wizard, StateView, Button
from trytond.pyson import Eval

__metaclass__ = PoolMeta
__all__ = [
    'Channel', 'PrestashopImportOrdersWizardView',
    'PrestashopImportOrdersWizard', 'PrestashopExportOrdersWizardView',
    'PrestashopExportOrdersWizard',
    'PrestashopConnectionWizardView', 'PrestashopConnectionWizard',
]
TIMEZONES = [(None, '')] + [(x, x) for x in pytz.common_timezones]

PRESTASHOP_STATES = {
    'required': Eval('source') == 'prestashop',
    'invisible': ~(Eval('source') == 'prestashop')
}
INVISIBLE_IF_NOT_PRESTASHOP = {
    'invisible': ~(Eval('source') == 'prestashop')
}

class Channel:
    """
    Sale Channel model
    """
    __name__ = 'sale.channel'

    #: The URL of prestashop site
    prestashop_url = fields.Char(
        'Prestashop URL', states=PRESTASHOP_STATES, depends=['source']
    )

    #: The webservices key for access to site
    prestashop_key = fields.Char(
        'Prestashop Key', states=PRESTASHOP_STATES, depends=['source']
    )

    #: Last time at which the orders were imported from prestashop
    last_prestashop_order_import_time = fields.DateTime(
        'Last Prestashop Order Import Time', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )

    #: Last time at which the orders were exported to prestashop
    last_prestashop_order_export_time = fields.DateTime(
        'Last Prestashop order export time', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )
    prestashop_shipping_product = fields.Many2One(
        'product.product', 'Shipping Product', states=PRESTASHOP_STATES,
        domain=[
            ('type', '=', 'service'),
            ('template.type', '=', 'service'),
            ('salable', '=', True),
        ], depends=['source']
    )

    #: The timezone set on prestashop site
    #: The orders imported will bear time in this timezone and that will need
    #: to be converted to UTC
    #: Also in order to determine what orders are to be imported, we need
    #: to convert UTC to this timezone to ensure correct time interval
    prestashop_timezone = fields.Selection(
        TIMEZONES, 'Timezone', translate=False, states=PRESTASHOP_STATES,
        depends=['source']
    )

    #: Allowed languages to be synced for the site
    prestashop_languages = fields.One2Many(
        'prestashop.site.lang', 'channel', 'Languages',
        states=INVISIBLE_IF_NOT_PRESTASHOP, depends=['source']
    )

    #: The mapping between prestashop order states and tryton sale states
    prestashop_order_states = fields.One2Many(
        'prestashop.site.order_state', 'channel', 'Order States',
        states=INVISIBLE_IF_NOT_PRESTASHOP, depends=['source']
    )

    #: Set this to True to handle invoicing in tryton and get invoice info
    #: TODO: A provision to be implemented in future versions
    #: Also handle multiple payment methods where each will have different
    #: journal and account setups
    prestashop_handle_invoice = fields.Boolean(
        'Handle Invoicing ?', states=INVISIBLE_IF_NOT_PRESTASHOP,
        depends=['source']
    )

    @classmethod
    def get_source(cls):
        """
        Get the source
        """
        res = super(Channel, cls).get_source()
        res.append(('prestashop', 'Prestashop'))
        return res

    def validate_prestashop_channel(self):
        """
        Check if current channel belongs to prestashop
        """
        if self.source != 'prestashop':
            self.raise_user_error("invalid_prestashop_channel")

    @classmethod
    def __setup__(cls):
        super(Channel, cls).__setup__()
        cls._error_messages.update({
            'prestashop_settings_missing':
                'Prestashop webservice settings are incomplete.',
            'invalid_prestashop_channel':
                'Current channel does not belongs to prestashop',
            'multiple_channels':
                'Test connection can be done for only one channel at a time.',
            'wrong_url_n_key': 'Connection Failed! Please check URL and Key',
            'wrong_url': 'Connection Failed! The URL provided is wrong',
            'languages_not_imported':
                'Import the languages before importing order states',
            'order_states_not_imported':
                'Import the order states before importing/exporting orders'
        })
        cls._buttons.update({
            'test_prestashop_connection': {},
            'import_prestashop_languages': {},
            'import_prestashop_order_states': {},
            'import_prestashop_orders': {},
            'export_prestashop_orders': {},
        })

    def get_prestashop_client(self):
        """
        Returns an authenticated instance of the Prestashop client

        :return: Prestashop client object
        """
        if not all([self.prestashop_url, self.prestashop_key]):
            self.raise_user_error('prestashop_settings_missing')

        if Transaction().context.get('ps_test'):
            return MockstaShopWebservice('Some URL', 'A Key')

        return pystashop.PrestaShopWebservice(
            self.prestashop_url, self.prestashop_key
        )

    @classmethod
    @ModelView.button
    def import_prestashop_languages(cls, channels):
        """Import Languages from remote and try to link them to tryton
        languages

        :param sites: List of sites for which the languages are to be imported
        :returns: List of languages created
        """
        SiteLanguage = Pool().get('prestashop.site.lang')

        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()

        # Set this site in context
        with Transaction().set_context(current_channel=channel.id):

            client = channel.get_prestashop_client()
            languages = client.languages.get_list(display='full')

            new_records = []
            for lang in languages:
                # If the language already exists in `Languages`, skip and do
                # not create it again
                if SiteLanguage.search_using_ps_id(lang.id.pyval):
                    continue
                new_records.append(
                    SiteLanguage.create_using_ps_data(lang)
                )

        return new_records

    @classmethod
    @ModelView.button
    def import_prestashop_order_states(cls, channels):
        """Import Order States from remote and try to link them to tryton
        order states

        :returns: List of order states created
        """
        SiteOrderState = Pool().get('prestashop.site.order_state')

        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()

        # Set this channel to context
        with Transaction().set_context(current_channel=channel.id):

            # If channel languages don't exist, then raise an error
            if not channel.prestashop_languages:
                cls.raise_user_error('languages_not_imported')

            client = channel.get_prestashop_client()
            order_states = client.order_states.get_list(display='full')

            new_records = []
            for state in order_states:
                # If this order state already exists for this channel, skip and
                # do not create it again
                if SiteOrderState.search_using_ps_id(state.id.pyval):
                    continue

                new_records.append(
                    SiteOrderState.create_using_ps_data(state)
                )

        return new_records

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_connection')
    def test_prestashop_connection(cls, channels):
        """Test Prestashop connection and display appropriate message to user
        """
        if len(channels) != 1:
            cls.raise_user_error('multiple_channels')
        channel = channels[0]

        channel.validate_prestashop_channel()
        try:
            client = channel.get_prestashop_client()

            # Try getting the list of shops
            # If it fails with prestashop error, then raise error
            client.shops.get_list()
        except pystashop.PrestaShopWebserviceException:
            cls.raise_user_error('wrong_url_n_key')
        except (
            requests.exceptions.MissingSchema,
            requests.exceptions.ConnectionError
        ):
            cls.raise_user_error('wrong_url')

    @classmethod
    def import_prestashop_orders_using_cron(cls, channels):
        """Import orders from prestashop

        :param sites: The list of sites from which the orders are to be
                      imported

        ..note:: This method is usually called by the cron
        """
        channels = cls.search([
            ('source', '=', 'prestashop')
        ])
        for channel in channels:
            channel.import_prestashop_orders()

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_import_orders')
    def import_prestashop_orders_button(cls, sites=None):
        """Dummy button to fire up the wizard for import of orders

        :param sites: The list of sites from which the orders are to be
                      imported
        """
        pass

    def import_prestashop_orders(self):
        """Import orders for the current site from prestashop
        Import only those orers which are updated after the
        `last order import time` as set in the prestashop configuration

        :returns: The list of active records of sales imported
        """
        Sale = Pool().get('sale.sale')
        self.validate_prestashop_channel()

        if not self.prestashop_order_states:
            self.raise_user_error('order_states_not_imported')

        # Localize to the site timezone
        utc_time_now = datetime.utcnow()
        site_tz = pytz.timezone(self.prestashop_timezone)
        time_now = site_tz.normalize(pytz.utc.localize(utc_time_now))
        client = self.get_prestashop_client()

        with Transaction().set_context(current_channel=self.id):
            if self.last_prestashop_order_import_time:
                # In tryton all time stored is in UTC
                # Convert the last import time to timezone of the site
                last_prestashop_order_import_time = site_tz.normalize(
                    pytz.utc.localize(self.last_prestashop_order_import_time)
                )
                orders_to_import = client.orders.get_list(
                    filters={
                        'date_upd': '{0},{1}'.format(
                            last_prestashop_order_import_time.strftime(
                                '%Y-%m-%d %H:%M:%S'
                            ),
                            time_now.strftime('%Y-%m-%d %H:%M:%S')
                        )
                    }, date=1, display='full'
                )
            else:
                # FIXME: This wont scale if there are thousands of orders
                orders_to_import = client.orders.get_list(display='full')

            self.write([self], {
                'last_prestashop_order_import_time': utc_time_now
            })
            sales_imported = []
            for order in orders_to_import:
                sales_imported.append(Sale.find_or_create_using_ps_data(order))

        return sales_imported

    @classmethod
    def export_orders_to_prestashop_using_cron(cls):
        """
        Export order status to prestashop using cron
        """
        channels = cls.search([
            ('source', '=', 'prestashop')
        ])
        for channel in channels:
            channel.export_orders_to_prestashop()

    @classmethod
    @ModelView.button_action('prestashop.wizard_prestashop_export_orders')
    def export_orders_button(cls, channels):
        """
        Dummy button to fire up the wizard for export of orders
        """
        pass

    def export_orders_to_prestashop(self):
        """
        Export order status to prestashop current site
        Export only those orders which are modified after the
        `last order export time` as set in the prestashop configuration.

        :returns: The list of active records of sales exported
        """
        Sale = Pool().get('sale.sale')
        Move = Pool().get('stock.move')

        if not self.order_states:
            self.raise_user_error('order_states_not_imported')

        time_now = datetime.utcnow()

        self.validate_prestashop_channel()

        with Transaction().set_context(current_channel=self.id):
            if self.last_prestashop_order_export_time:
                # Sale might not get updated for state changes in the related
                # shipments.
                # So first get the moves for outgoing shipments which are \
                # executed after last import time.
                moves = Move.search([
                    (
                        'write_date', '>=',
                        self.last_prestashop_order_export_time
                    ),
                    ('sale.channel', '=', self.id),
                    ('shipment', 'like', 'stock.shipment.out%')
                ])
                sales_to_export = Sale.search(['OR', [
                    ('write_date', '>=', self.last_order_export_time),
                    ('channel', '=', self.id),
                ], [
                    ('id', 'in', map(int, [m.sale for m in moves]))
                ]])
            else:
                sales_to_export = Sale.search([('channel', '=', self.id)])

            self.write([self], {
                'last_prestashop_order_export_time': time_now
            })

            for sale in sales_to_export:
                sale.export_status_to_ps()

        return sales_to_export

class PrestashopConnectionWizardView(ModelView):
    'Prestashop Connection Wizard View'
    __name__ = 'prestashop.connection.wizard.view'

class PrestashopConnectionWizard(Wizard):
    'Prestashop Connection Wizard'
    __name__ = 'prestashop.connection.wizard'

    start = StateView(
        'prestashop.connection.wizard.view',
        'prestashop.prestashop_connection_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, data):
        """Test the connection and show the user appropriate message

        :param data: Wizard data
        """
        return {}

class PrestashopImportOrdersWizardView(ModelView):
    'Prestashop Import Wizard View'
    __name__ = 'prestashop.import_orders.wizard.view'

    orders_imported = fields.Integer('Orders Imported', readonly=True)

class PrestashopImportOrdersWizard(Wizard):
    'Prestashop Import Wizard'
    __name__ = 'prestashop.import_orders.wizard'

    start = StateView(
        'prestashop.import_orders.wizard.view',
        'prestashop.prestashop_import_orders_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, fields):
        """
        Import the orders and display a confirmation message to the user

        :param fields: Wizard fields
        """
        SaleChannel = Pool().get('sale.channel')

        channel = SaleChannel(Transaction().context['active_id'])

        channel.validate_prestashop_channel()

        return {
            'orders_imported': len(channel.import_prestashop_orders())
        }

class PrestashopExportOrdersWizardView(ModelView):
    'Prestashop Export Orders Wizard View'
    __name__ = 'prestashop.export_orders.wizard.view'

    orders_exported = fields.Integer('Orders Exported', readonly=True)

class PrestashopExportOrdersWizard(Wizard):
    'Prestashop Export Orders Wizard'
    __name__ = 'prestashop.export_orders.wizard'

    start = StateView(
        'prestashop.export_orders.wizard.view',
        'prestashop.prestashop_export_orders_wizard_view_form',
        [
            Button('Ok', 'end', 'tryton-ok'),
        ]
    )

    def default_start(self, fields):
        """
        Export the orders and display a confirmation message to the user

        :param fields: Wizard fields
        """
        SaleChannel = Pool().get('sale.channel')

        channel = SaleChannel(Transaction().context['active_id'])

        channel.validate_prestashop_channel()

        default = {
            'orders_exported': len(channel.export_orders_to_prestashop())
        }

        return default
<|endoftext|>


<commit_before>#!/usr/bin/env python

import os
import sys
import gzip
import re
import subprocess
import datetime
import random
import traceback

from optparse import OptionParser

try:
	import simplejson as json
except ImportError:
	import json

parent = os.path.dirname
basename = os.path.basename
join = os.path.join

BAD_FNAME_RE = re.compile(r"""[\x00"'\\]+""")

def check_filename(fname):
	"""
	Raise L{ValueError} on any filename that can't be safely passed into a quoted
	shell command.
	"""
	if re.findall(BAD_FNAME_RE, fname):
		raise ValueError("Bad filename %r" % (fname,))

def try_makedirs(p):
	try:
		os.makedirs(p)
	except OSError:
		pass

def slurp_gz(fname):
	f = gzip.open(fname, "rb")
	try:
		contents = f.read()
	finally:
		f.close()
	return contents

def get_expected_encoded_feed_urls(greader_items, item_name):
	return slurp_gz(join(greader_items, item_name[0:6], item_name + '.gz')).rstrip("\n").split("\n")

def full_greader_url(encoded_feed_url):
	return (
		"https://www.google.com/reader/api/0/stream/contents/feed/" +
		  encoded_feed_url +
		"?r=n&n=1000&hl=en&likes=true&comments=true&client=ArchiveTeam")

CONTINUATION_RE = re.compile(r"\?c=............&")
QUESTION_RE = re.compile(r"\?")

def url_with_continuation(url, continuation):
	assert len(continuation) == 12, "continuation should be 12 bytes, was %d" % (len(continuation),)
	if re.findall(CONTINUATION_RE, url):
		return re.sub(CONTINUATION_RE, "?c=" + continuation + "&", url)
	return re.sub(QUESTION_RE, "?c=" + continuation + "&", url, count=1)

def is_continued_url(url):
	return '?c=' in url

class BadWARC(Exception):
	pass

def read_request_responses(grepfh, hrefs):
	"""
	L{grepfh} is the file object containing grep-filtered WARC data.

	L{hrefs} is a set into which this will add JSON-encoded hrefs to.
	"""
	WANT_FIRST_TARGET_URI, NEED_SECOND_TARGET_URL, NEED_STATUS_LINE, WANT_CONTINUATION = range(4)
	state = WANT_FIRST_TARGET_URI
	last_url = None
	continuation = None
	status_code = None
	while True:
		line = grepfh.readline()
		if not line:
			if last_url is not None:
				if status_code is None:
					raise BadWARC("Did not get a status code for %r" % (last_url,))
				yield dict(url=last_url, continuation=continuation, status_code=status_code)
			break

		if state == WANT_FIRST_TARGET_URI:
			# in this state, we may get the Target-URI, or we may get an href
			if line.startswith("WARC-Target-URI: "):
				if last_url is not None:
					if status_code is None:
						raise BadWARC("Did not get a status code for %r" % (last_url,))
					yield dict(url=last_url, continuation=continuation, status_code=status_code)
				continuation = None
				status_code = None
				last_url = line[17:-2]
				if last_url[:4] != "http": # skip warc metadata
					last_url = None
					state = WANT_FIRST_TARGET_URI
				else:
					state = NEED_SECOND_TARGET_URL
			elif line.startswith(r'href\u003d\"'):
				hrefs.add(line[12:-3])
			else:
				# Ignore unexpected lines, as the lack of ^ in our initial grep filter
				# outputs some garbage
				pass

		elif state == NEED_SECOND_TARGET_URL:
			# Should be an exact duplicate of the last line
			assert last_url is not None, last_url
			if not line.startswith("WARC-Target-URI: "):
				raise BadWARC("Missing WARC-Target-URI for response")
			response_url = line[17:-2]
			if response_url != last_url:
				raise BadWARC("WARC-Target-URI for response did not match request: %r" % ((last_url, response_url),))
			state = NEED_STATUS_LINE

		elif state == NEED_STATUS_LINE:
			try:
				http_version, status_code, message = line.split(" ", 2)
			except ValueError:
				raise BadWARC("Got unexpected status line %r" % (line,))
			if http_version != "HTTP/1.1"or status_code not in ("200", "404"):
				raise BadWARC("Got unexpected status line %r" % (line,))
			state = WANT_CONTINUATION

		elif state == WANT_CONTINUATION:
			# could get continuation (once chance at this), or a link, or next request
			if line.startswith('"continuation":"'):
				continuation = line[16:28]
				state = WANT_FIRST_TARGET_URI
			elif line.startswith(r'href\u003d\"'):
				hrefs.add(line[12:-3])
				state = WANT_FIRST_TARGET_URI
			elif line.startswith("WARC-Target-URI: "):
				yield dict(url=last_url, continuation=continuation, status_code=status_code)
				continuation = None
				status_code = None
				last_url = line[17:-2]
				if last_url[:4] != "http": # skip warc metadata
					last_url = None
					state = WANT_FIRST_TARGET_URI
				else:
					state = NEED_SECOND_TARGET_URL
			else:
				# Ignore unexpected lines, as the lack of ^ in our initial grep filter
				# outputs some garbage
				pass

		else:
			raise RuntimeError("Invalid state %r" % (state,))

def get_info_from_warc_fname(fname):
	"""
	Where fname is absolute path, or at least includes the uploader parent dir
	"""
	uploader = basename(parent(fname))
	_, item_name, _, _ = basename(fname).split('-')
	return dict(uploader=uploader, item_name=item_name, basename=basename(fname))

def check_warc(fname, info, greader_items, href_log, reqres_log):
	uploader = info['uploader']
	item_name = info['item_name']
	expected_urls = set(full_greader_url(efu) for efu in get_expected_encoded_feed_urls(greader_items, item_name))

	check_filename(fname)

	# We use pipes to allow for multi-core execution without writing a crazy amount
	# of Python code that wires up subprocesses.
	# "Z8c8Jv5QWmpgVRxUsGoulMw" is the embedded 404 image we want to ignore.
	# Do not add a ^ to the second grep - it will slow things 6x.
	args = ['/bin/sh', '-c', r"""gunzip --to-stdout '%s' | grep -G --color=never -v "^Z8c8Jv5QWmpgVRxUsGoulMw" | grep -P --color=never -o 'href\\u003d\\"[^\\]+\\"|"continuation":"C.{10}C"|WARC-Target-URI: .*|HTTP/1\.1 .*'""" % (fname,)]
	proc = subprocess.Popen(args, stdout=subprocess.PIPE)
	found_hrefs = set()
	got_urls = set()
	# Note: if gzip file is corrupt, stdout will be empty and a BadWARC will be raised
	# complaining about the WARC missing every URL it was expected to have.
	for req_rep in read_request_responses(proc.stdout, found_hrefs):
		req_rep_extra = dict(item_name=item_name, uploader=uploader, basename=info['basename'], **req_rep)
		##print json.dumps(req_rep_extra)
		json.dump(req_rep_extra, reqres_log)
		reqres_log.write("\n")
		url = req_rep['url']
		status_code = req_rep['status_code']
		if req_rep['continuation'] is not None:
			expected_urls.add(url_with_continuation(url, req_rep['continuation']))
		got_urls.add(url)
		if is_continued_url(url) and status_code != "200":
			raise BadWARC("All continued responses must be status 200, was %r" % (status_code,))

	if expected_urls != got_urls:
		raise BadWARC("WARC is missing %r or has extra URLs %r" % (expected_urls - got_urls, got_urls - expected_urls))

	sorted_hrefs = list(found_hrefs)
	sorted_hrefs.sort()
	if href_log is not None:
		for h in sorted_hrefs:
			href_log.write(h + "\n")

def main():
	parser = OptionParser(usage="%prog [options]")

	parser.add_option("-i", "--input-base", dest="input_base", help="Base directory containing ./username/xxx.warc.gz files.")
	parser.add_option("-o", "--output-base", dest="output_base", help="Base directory to which to move input files; it will contain ./verified/username/xxx.warc.gz or ./bad-[failure mode]/username/xxx.warc.gz.  Should be on the same filesystem as --input-base.")
	parser.add_option('-g', "--greader-items", dest="greader_items", help="greader-items directory containing ./000000/0000000000.gz files.  (Needed to know which URLs we expect in a WARC.)")
	parser.add_option("-l", "--lists-dir", dest="lists_dir", help="Directory to write lists of status codes, bad items, new URLs to.")
	parser.add_option("-u", "--upload", dest="upload", help="rsync destination to sync lists to.")

	options, args = parser.parse_args()
	if not options.input_base or  not options.greader_items:
		print "--input-base and --greader-items are required"
		print
		parser.print_help()
		sys.exit(1)

	if not options.output_base:
		print "--output-base not specified; files in --input-base will not be moved"

	if not options.lists_dir:
		print "--lists-dir not specified; no lists will be written"

	if not options.upload:
		print "--upload not specified; lists will not be uploaded"

	if options.output_base:
		verified_dir = join(options.output_base, "verified")
		bad_dir = join(options.output_base, "bad")
	else:
		verified_dir = bad_dir = None

	if options.lists_dir:
		now = datetime.datetime.now()
		full_date = now.isoformat().replace("T", "_").replace(':', '-') + "_" + str(random.random())[2:8]

		href_log_fname = join(options.lists_dir, full_date + ".hrefs")
		assert not os.path.exists(href_log_fname), href_log_fname
		href_log = open(href_log_fname, "wb")

		reqres_log_fname = join(options.lists_dir, full_date + ".reqres")
		assert not os.path.exists(reqres_log_fname), reqres_log_fname
		reqres_log = open(reqres_log_fname, "wb")

		verification_log_fname = join(options.lists_dir, full_date + ".verification")
		assert not os.path.exists(verification_log_fname), verification_log_fname
		verification_log = open(verification_log_fname, "wb")
	else:
		href_log = None
		reqres_log = None
		verification_log = None

	for directory, dirnames, filenames in os.walk(options.input_base):
		for f in filenames:
			fname = os.path.join(directory, f)
			if fname.endswith('.warc.gz'):
				info = get_info_from_warc_fname(fname)
				try:
					check_warc(fname, info, options.greader_items, href_log, reqres_log)
				except BadWARC, e:
					# TODO move the file to bad/ instead
					json.dump(dict(valid=False, exception=repr(e), traceback=traceback.format_exc(), **info), verification_log)
					verification_log.write("\n")
					verification_log.flush()
				else:
					json.dump(dict(valid=True, exception=None, traceback=None, **info), verification_log)
					verification_log.write("\n")
					verification_log.flush()

if __name__ == '__main__':
	main()
<commit_msg>No need for exception=; traceback has everything<commit_after>#!/usr/bin/env python

import os
import sys
import gzip
import re
import subprocess
import datetime
import random
import traceback

from optparse import OptionParser

try:
	import simplejson as json
except ImportError:
	import json

parent = os.path.dirname
basename = os.path.basename
join = os.path.join

BAD_FNAME_RE = re.compile(r"""[\x00"'\\]+""")

def check_filename(fname):
	"""
	Raise L{ValueError} on any filename that can't be safely passed into a quoted
	shell command.
	"""
	if re.findall(BAD_FNAME_RE, fname):
		raise ValueError("Bad filename %r" % (fname,))

def try_makedirs(p):
	try:
		os.makedirs(p)
	except OSError:
		pass

def slurp_gz(fname):
	f = gzip.open(fname, "rb")
	try:
		contents = f.read()
	finally:
		f.close()
	return contents

def get_expected_encoded_feed_urls(greader_items, item_name):
	return slurp_gz(join(greader_items, item_name[0:6], item_name + '.gz')).rstrip("\n").split("\n")

def full_greader_url(encoded_feed_url):
	return (
		"https://www.google.com/reader/api/0/stream/contents/feed/" +
		  encoded_feed_url +
		"?r=n&n=1000&hl=en&likes=true&comments=true&client=ArchiveTeam")

CONTINUATION_RE = re.compile(r"\?c=............&")
QUESTION_RE = re.compile(r"\?")

def url_with_continuation(url, continuation):
	assert len(continuation) == 12, "continuation should be 12 bytes, was %d" % (len(continuation),)
	if re.findall(CONTINUATION_RE, url):
		return re.sub(CONTINUATION_RE, "?c=" + continuation + "&", url)
	return re.sub(QUESTION_RE, "?c=" + continuation + "&", url, count=1)

def is_continued_url(url):
	return '?c=' in url

class BadWARC(Exception):
	pass

def read_request_responses(grepfh, hrefs):
	"""
	L{grepfh} is the file object containing grep-filtered WARC data.

	L{hrefs} is a set into which this will add JSON-encoded hrefs to.
	"""
	WANT_FIRST_TARGET_URI, NEED_SECOND_TARGET_URL, NEED_STATUS_LINE, WANT_CONTINUATION = range(4)
	state = WANT_FIRST_TARGET_URI
	last_url = None
	continuation = None
	status_code = None
	while True:
		line = grepfh.readline()
		if not line:
			if last_url is not None:
				if status_code is None:
					raise BadWARC("Did not get a status code for %r" % (last_url,))
				yield dict(url=last_url, continuation=continuation, status_code=status_code)
			break

		if state == WANT_FIRST_TARGET_URI:
			# in this state, we may get the Target-URI, or we may get an href
			if line.startswith("WARC-Target-URI: "):
				if last_url is not None:
					if status_code is None:
						raise BadWARC("Did not get a status code for %r" % (last_url,))
					yield dict(url=last_url, continuation=continuation, status_code=status_code)
				continuation = None
				status_code = None
				last_url = line[17:-2]
				if last_url[:4] != "http": # skip warc metadata
					last_url = None
					state = WANT_FIRST_TARGET_URI
				else:
					state = NEED_SECOND_TARGET_URL
			elif line.startswith(r'href\u003d\"'):
				hrefs.add(line[12:-3])
			else:
				# Ignore unexpected lines, as the lack of ^ in our initial grep filter
				# outputs some garbage
				pass

		elif state == NEED_SECOND_TARGET_URL:
			# Should be an exact duplicate of the last line
			assert last_url is not None, last_url
			if not line.startswith("WARC-Target-URI: "):
				raise BadWARC("Missing WARC-Target-URI for response")
			response_url = line[17:-2]
			if response_url != last_url:
				raise BadWARC("WARC-Target-URI for response did not match request: %r" % ((last_url, response_url),))
			state = NEED_STATUS_LINE

		elif state == NEED_STATUS_LINE:
			try:
				http_version, status_code, message = line.split(" ", 2)
			except ValueError:
				raise BadWARC("Got unexpected status line %r" % (line,))
			if http_version != "HTTP/1.1"or status_code not in ("200", "404"):
				raise BadWARC("Got unexpected status line %r" % (line,))
			state = WANT_CONTINUATION

		elif state == WANT_CONTINUATION:
			# could get continuation (once chance at this), or a link, or next request
			if line.startswith('"continuation":"'):
				continuation = line[16:28]
				state = WANT_FIRST_TARGET_URI
			elif line.startswith(r'href\u003d\"'):
				hrefs.add(line[12:-3])
				state = WANT_FIRST_TARGET_URI
			elif line.startswith("WARC-Target-URI: "):
				yield dict(url=last_url, continuation=continuation, status_code=status_code)
				continuation = None
				status_code = None
				last_url = line[17:-2]
				if last_url[:4] != "http": # skip warc metadata
					last_url = None
					state = WANT_FIRST_TARGET_URI
				else:
					state = NEED_SECOND_TARGET_URL
			else:
				# Ignore unexpected lines, as the lack of ^ in our initial grep filter
				# outputs some garbage
				pass

		else:
			raise RuntimeError("Invalid state %r" % (state,))

def get_info_from_warc_fname(fname):
	"""
	Where fname is absolute path, or at least includes the uploader parent dir
	"""
	uploader = basename(parent(fname))
	_, item_name, _, _ = basename(fname).split('-')
	return dict(uploader=uploader, item_name=item_name, basename=basename(fname))

def check_warc(fname, info, greader_items, href_log, reqres_log):
	uploader = info['uploader']
	item_name = info['item_name']
	expected_urls = set(full_greader_url(efu) for efu in get_expected_encoded_feed_urls(greader_items, item_name))

	check_filename(fname)

	# We use pipes to allow for multi-core execution without writing a crazy amount
	# of Python code that wires up subprocesses.
	# "Z8c8Jv5QWmpgVRxUsGoulMw" is the embedded 404 image we want to ignore.
	# Do not add a ^ to the second grep - it will slow things 6x.
	args = ['/bin/sh', '-c', r"""gunzip --to-stdout '%s' | grep -G --color=never -v "^Z8c8Jv5QWmpgVRxUsGoulMw" | grep -P --color=never -o 'href\\u003d\\"[^\\]+\\"|"continuation":"C.{10}C"|WARC-Target-URI: .*|HTTP/1\.1 .*'""" % (fname,)]
	proc = subprocess.Popen(args, stdout=subprocess.PIPE)
	found_hrefs = set()
	got_urls = set()
	# Note: if gzip file is corrupt, stdout will be empty and a BadWARC will be raised
	# complaining about the WARC missing every URL it was expected to have.
	for req_rep in read_request_responses(proc.stdout, found_hrefs):
		req_rep_extra = dict(item_name=item_name, uploader=uploader, basename=info['basename'], **req_rep)
		##print json.dumps(req_rep_extra)
		json.dump(req_rep_extra, reqres_log)
		reqres_log.write("\n")
		url = req_rep['url']
		status_code = req_rep['status_code']
		if req_rep['continuation'] is not None:
			expected_urls.add(url_with_continuation(url, req_rep['continuation']))
		got_urls.add(url)
		if is_continued_url(url) and status_code != "200":
			raise BadWARC("All continued responses must be status 200, was %r" % (status_code,))

	if expected_urls != got_urls:
		raise BadWARC("WARC is missing %r or has extra URLs %r" % (expected_urls - got_urls, got_urls - expected_urls))

	sorted_hrefs = list(found_hrefs)
	sorted_hrefs.sort()
	if href_log is not None:
		for h in sorted_hrefs:
			href_log.write(h + "\n")

def main():
	parser = OptionParser(usage="%prog [options]")

	parser.add_option("-i", "--input-base", dest="input_base", help="Base directory containing ./username/xxx.warc.gz files.")
	parser.add_option("-o", "--output-base", dest="output_base", help="Base directory to which to move input files; it will contain ./verified/username/xxx.warc.gz or ./bad-[failure mode]/username/xxx.warc.gz.  Should be on the same filesystem as --input-base.")
	parser.add_option('-g', "--greader-items", dest="greader_items", help="greader-items directory containing ./000000/0000000000.gz files.  (Needed to know which URLs we expect in a WARC.)")
	parser.add_option("-l", "--lists-dir", dest="lists_dir", help="Directory to write lists of status codes, bad items, new URLs to.")
	parser.add_option("-u", "--upload", dest="upload", help="rsync destination to sync lists to.")

	options, args = parser.parse_args()
	if not options.input_base or  not options.greader_items:
		print "--input-base and --greader-items are required"
		print
		parser.print_help()
		sys.exit(1)

	if not options.output_base:
		print "--output-base not specified; files in --input-base will not be moved"

	if not options.lists_dir:
		print "--lists-dir not specified; no lists will be written"

	if not options.upload:
		print "--upload not specified; lists will not be uploaded"

	if options.output_base:
		verified_dir = join(options.output_base, "verified")
		bad_dir = join(options.output_base, "bad")
	else:
		verified_dir = bad_dir = None

	if options.lists_dir:
		now = datetime.datetime.now()
		full_date = now.isoformat().replace("T", "_").replace(':', '-') + "_" + str(random.random())[2:8]

		href_log_fname = join(options.lists_dir, full_date + ".hrefs")
		assert not os.path.exists(href_log_fname), href_log_fname
		href_log = open(href_log_fname, "wb")

		reqres_log_fname = join(options.lists_dir, full_date + ".reqres")
		assert not os.path.exists(reqres_log_fname), reqres_log_fname
		reqres_log = open(reqres_log_fname, "wb")

		verification_log_fname = join(options.lists_dir, full_date + ".verification")
		assert not os.path.exists(verification_log_fname), verification_log_fname
		verification_log = open(verification_log_fname, "wb")
	else:
		href_log = None
		reqres_log = None
		verification_log = None

	for directory, dirnames, filenames in os.walk(options.input_base):
		for f in filenames:
			fname = os.path.join(directory, f)
			if fname.endswith('.warc.gz'):
				info = get_info_from_warc_fname(fname)
				try:
					check_warc(fname, info, options.greader_items, href_log, reqres_log)
				except BadWARC, e:
					# TODO move the file to bad/ instead
					json.dump(dict(valid=False, traceback=traceback.format_exc(), **info), verification_log)
					verification_log.write("\n")
					verification_log.flush()
				else:
					json.dump(dict(valid=True, traceback=None, **info), verification_log)
					verification_log.write("\n")
					verification_log.flush()

if __name__ == '__main__':
	main()
<|endoftext|>


<commit_before>#
# Copyright (c) 2004 Specifix, Inc.
# All rights reserved
#

import changeset
import cook
import helper
import log
import os
import package
import recipe
import sys
import update
import versioned
import versions

class SourceState(package.Package):

    def removeFilePath(self, file):
	for (fileId, (path, version)) in self.iterFileList():
	    if path == file:
		self.removeFile(fileId)
		return True

	return False

    def write(self, filename):
	f = open(filename, "w")
	f.write("name %s\n" % self.name)
	if self.version:
	    f.write("version %s\n" % self.version.asString())
	f.write("branch %s\n" % self.branch.asString())
	f.write(self.formatString())

    def getBranch(self):
	return self.branch

    def changeBranch(self, branch):
	self.branch = branch

    def getRecipeFileName(self):
        # XXX this is not the correct way to solve this problem
        # assumes a fully qualified trove name
        fields = self.getName().split(':')
        name = fields[-2]
        return os.path.join(os.getcwd(), name + '.recipe')

    def expandVersionStr(self, versionStr):
	if versionStr[0] == "@":
	    # get the name of the repository from the current branch
	    repName = self.getBranch().branchNickname().getHost()
	    return repName + versionStr
	elif versionStr[0] != "/" and versionStr.find("@") == -1:
	    # non fully-qualified version; make it relative to the current
	    # branch
	    return self.getBranch().asString() + "/" + versionStr

    def __init__(self, name, version, branch):
	package.Package.__init__(self, name, version)
	self.branch = branch

class SourceStateFromFile(SourceState):

    def parseFile(self, filename):
	f = open(filename)
	rc = [self]
	for (what, isBranch) in [ ('name', 0), ('version', 1), ('branch', 1) ]:
	    line = f.readline()
	    fields = line.split()
	    assert(len(fields) == 2)
	    assert(fields[0] == what)
	    if isBranch:
		rc.append(versions.VersionFromString(fields[1]))
	    else:
		rc.append(fields[1])

	SourceState.__init__(*rc)

	self.read(f)

    def __init__(self, file):
	if not os.path.isfile(file):
	    log.error("SRS file must exist in the current directory for source commands")
	    raise OSError  # XXX

	self.parseFile(file)

def _verifyAtHead(repos, headPkg, state):
    headVersion = repos.pkgLatestVersion(state.getName(),
					 state.getBranch())
    if not headVersion.equal(state.getVersion()):
	return False

    # make sure the files in this directory are based on the same
    # versions as those in the package at head
    bail = 0
    for (fileId, (path, version)) in state.iterFileList():
	if isinstance(version, versions.NewVersion):
	    assert(not headPkg.hasFile(fileId))
	    # new file, it shouldn't be in the old package at all
	else:
	    srcFileVersion = headPkg.getFile(fileId)[1]
	    if not version.equal(srcFileVersion):
		return False

    return True

def _getRecipeVersion(recipeFile):
    # load the recipe; we need this to figure out what version we're building
    try:
        loader = recipe.RecipeLoader(recipeFile)
    except recipe.RecipeFileError, e:
	log.error("unable to load recipe file %s: %s",
                  state.getRecipeFileName(), str(e))
        return None

    if not loader:
	log.error("unable to load a valid recipe class from %s",
                  state.getRecipeFileName())
	return None

    assert(len(loader.values()) == 1)
    recipeClass = loader.values()[0]

    return recipeClass.version

def checkout(repos, cfg, dir, name, versionStr = None):
    # We have to be careful with branch nicknames.  First, we could get
    # multiple matches for a single package. Two, when a nickname uniquely
    # identifies a package we still need to make sure the state has the name of
    # the actual branch since empty branches yield objects whose versions are
    # on the parent branch.
    name += ":sources"
    try:
        trvList = helper.findPackage(repos, cfg.packagenamespace,
                                     cfg.installbranch, name,
                                     versionStr = versionStr)
    except helper.PackageNotFound, e:
        log.error(str(e))
        return
    if len(trvList) > 1:
	log.error("branch %s matches more then one version", versionStr)
	return
    trv = trvList[0]

    if not dir:
	dir = trv.getName().split(":")[-2]

    if not os.path.isdir(dir):
	try:
	    os.mkdir(dir)
	except OSError, err:
	    log.error("cannot create directory %s/%s: %s", os.getcwd(), dir,
                      str(err))
	    return

    branch = helper.fullBranchName(cfg.packagenamespace[1:], cfg.installbranch,
				   trv.getVersion(), versionStr)
    state = SourceState(trv.getName(), trv.getVersion(), branch)

    for (fileId, (path, version)) in trv.iterFileList():
	fullPath = dir + "/" + path
	(fileObj, contents) = repos.getFileVersion(fileId, version,
						   withContents = True)

	fileObj.restore(contents, fullPath, 1)

	state.addFile(fileId, path, version)

    state.write(dir + "/SRS")

def commit(repos):
    state = SourceStateFromFile("SRS")

    if isinstance(state.getVersion(), versions.NewVersion):
	# new package, so it shouldn't exist yet
	if repos.hasPackage(state.getName()):
	    log.error("%s is marked as a new package but it "
		      "already exists" % state.getName())
	    return
	srcPkg = None
    else:
	srcPkg = repos.getPackageVersion(state.getName(), state.getVersion())
	# update the version to one w/ a timestamp
	srcPkg.changeVersion(repos.pkgGetFullVersion(state.getName(),
			     srcPkg.getVersion()))

	if not _verifyAtHead(repos, srcPkg, state):
	    log.error("contents of working directory are not all "
		      "from the head of the branch; use update")
	    return

    recipeVersionStr = _getRecipeVersion(state.getRecipeFileName())
    if not recipeVersionStr: return

    if srcPkg:
	newVersion = helper.nextVersion(recipeVersionStr, srcPkg.getVersion(),
					state.getBranch(), binary = False)
    else:
	newVersion = helper.nextVersion(recipeVersionStr, None,
					state.getBranch(), binary = False)

    result = update.buildLocalChanges(repos, [(state, srcPkg, newVersion)])
    if not result: return

    (changeSet, ((isDifferent, newState),)) = result

    if not isDifferent:
	log.info("no changes have been made to commit")
    else:
	repos.commitChangeSet(changeSet)
	newState.write("SRS")

def diff(repos, versionStr = None):
    state = SourceStateFromFile("SRS")

    if versionStr:
	versionStr = state.expandVersionStr(versionStr)

	pkgList = helper.findPackage(repos, None, None, state.getName(),
				     versionStr)
	if len(pkgList) > 1:
	    log.error("%s specifies multiple versions" % versionStr)
	    return

	oldPackage = pkgList[0]
    else:
	oldPackage = repos.getPackageVersion(state.getName(),
					     state.getVersion())

    # update the version to one w/ a timestamp
    oldPackage.changeVersion(repos.pkgGetFullVersion(state.getName(),
			     oldPackage.getVersion()))

    result = update.buildLocalChanges(repos, [(state, oldPackage,
					       versions.NewVersion())])
    if not result: return

    (changeSet, ((isDifferent, newState),)) = result
    if not isDifferent: return

    packageChanges = changeSet.getNewPackageList()
    assert(len(packageChanges) == 1)
    pkgCs = packageChanges[0]

    for (fileId, path, newVersion) in pkgCs.getNewFileList():
	print "%s: new" % path

    for (fileId, path, newVersion) in pkgCs.getChangedFileList():
	if not path:
	    path = oldPackage.getFile(fileId)[0]
	    sys.stdout.write("%s" % path)
	else:
	    oldPath = oldPackage.getFile(fileId)[0]
	    sys.stdout.write("%s (aka %s)" % (path, oldPath))

	if not newVersion:
	    print
	    continue

	sys.stdout.write(": changed\n")

	csInfo = changeSet.getFileChange(fileId)
	print "    %s" % csInfo

	fileType = csInfo.split()[0]
	sha1 = csInfo.split()[1]
	if fileType == "f" and sha1 != "-":
	    (contType, contents) = changeSet.getFileContents(fileId)
	    if contType == changeset.ChangedFileTypes.diff:
		lines = contents.get().readlines()
		str = "    " + "    ".join(lines)
		print
		print str
		print

    for fileId in pkgCs.getOldFileList():
	path = oldPackage.getFile(fileId)[0]
	print "%s: removed" % path

def updateSrc(repos, versionStr = None):
    state = SourceStateFromFile("SRS")
    pkgName = state.getName()
    baseVersion = state.getVersion()

    if not versionStr:
	head = repos.getLatestPackage(pkgName, state.getBranch())
	newBranch = None
	headVersion = head.getVersion()
	if headVersion.equal(baseVersion):
	    log.info("working directory is already based on head of branch")
	    return
    else:
	versionStr = state.expandVersionStr(versionStr)

	pkgList = helper.findPackage(repos, None, None, pkgName, versionStr)
	if len(pkgList) > 1:
	    log.error("%s specifies multiple versions" % versionStr)
	    return

	head = pkgList[0]
	headVersion = head.getVersion()
	newBranch = helper.fullBranchName(None, None, headVersion, versionStr)

    changeSet = repos.createChangeSet([(pkgName, baseVersion, headVersion, 0)])

    packageChanges = changeSet.getNewPackageList()
    assert(len(packageChanges) == 1)
    pkgCs = packageChanges[0]

    fsJob = update.FilesystemJob(repos, changeSet,
				 { state.getName() : state }, "" )
    errList = fsJob.getErrorList()
    if errList:
	for err in errList: log.error(err)
    fsJob.apply()
    newPkgs = fsJob.getNewPackageList()
    assert(len(newPkgs) == 1)
    newState = newPkgs[0]

    if newState.getVersion().equal(pkgCs.getNewVersion()) and newBranch:
	newState.changeBranch(newBranch)

    newState.write("SRS")

def addFile(file):
    state = SourceStateFromFile("SRS")

    try:
	os.lstat(file)
    except OSError:
	log.error("files must be created before they can be added")
	return

    for (fileId, (path, version)) in state.iterFileList():
	if path == file:
	    log.error("file %s is already part of this source package" % path)
	    return

    fileId = cook.makeFileId(os.getcwd(), file)

    state.addFile(fileId, file, versions.NewVersion())
    state.write("SRS")

def removeFile(file):
    state = SourceStateFromFile("SRS")

    if not state.removeFilePath(file):
	log.error("file %s is not under management" % file)

    if os.path.exists(file):
	os.unlink(file)

    state.write("SRS")

def newPackage(repos, cfg, name):
    if name[0] != ":":
	name = cfg.packagenamespace + ":" + name
    name += ":sources"

    state = SourceState(name, versions.NewVersion(), cfg.defaultbranch)

    if repos and repos.hasPackage(name):
	log.error("package %s already exists" % name)
	return

    dir = name.split(":")[-2]
    if not os.path.isdir(dir):
	try:
	    os.mkdir(dir)
	except:
	    log.error("cannot create directory %s/%s", os.getcwd(), dir)
	    return

    state.write(dir + "/" + "SRS")

def renameFile(oldName, newName):
    state = SourceStateFromFile("SRS")

    if not os.path.exists(oldName):
	log.error("%s does not exist or is not a regular file" % oldName)
	return

    try:
	os.lstat(newName)
    except:
	pass
    else:
	log.error("%s already exists" % newName)
	return

    for (fileId, (path, version)) in state.iterFileList():
	if path == oldName:
	    os.rename(oldName, newName)
	    state.addFile(fileId, newName, version)
	    state.write("SRS")
	    return

    log.error("file %s is not under management" % oldName)
<commit_msg>fix error case, use recipeFile now<commit_after>#
# Copyright (c) 2004 Specifix, Inc.
# All rights reserved
#

import changeset
import cook
import helper
import log
import os
import package
import recipe
import sys
import update
import versioned
import versions

class SourceState(package.Package):

    def removeFilePath(self, file):
	for (fileId, (path, version)) in self.iterFileList():
	    if path == file:
		self.removeFile(fileId)
		return True

	return False

    def write(self, filename):
	f = open(filename, "w")
	f.write("name %s\n" % self.name)
	if self.version:
	    f.write("version %s\n" % self.version.asString())
	f.write("branch %s\n" % self.branch.asString())
	f.write(self.formatString())

    def getBranch(self):
	return self.branch

    def changeBranch(self, branch):
	self.branch = branch

    def getRecipeFileName(self):
        # XXX this is not the correct way to solve this problem
        # assumes a fully qualified trove name
        fields = self.getName().split(':')
        name = fields[-2]
        return os.path.join(os.getcwd(), name + '.recipe')

    def expandVersionStr(self, versionStr):
	if versionStr[0] == "@":
	    # get the name of the repository from the current branch
	    repName = self.getBranch().branchNickname().getHost()
	    return repName + versionStr
	elif versionStr[0] != "/" and versionStr.find("@") == -1:
	    # non fully-qualified version; make it relative to the current
	    # branch
	    return self.getBranch().asString() + "/" + versionStr

    def __init__(self, name, version, branch):
	package.Package.__init__(self, name, version)
	self.branch = branch

class SourceStateFromFile(SourceState):

    def parseFile(self, filename):
	f = open(filename)
	rc = [self]
	for (what, isBranch) in [ ('name', 0), ('version', 1), ('branch', 1) ]:
	    line = f.readline()
	    fields = line.split()
	    assert(len(fields) == 2)
	    assert(fields[0] == what)
	    if isBranch:
		rc.append(versions.VersionFromString(fields[1]))
	    else:
		rc.append(fields[1])

	SourceState.__init__(*rc)

	self.read(f)

    def __init__(self, file):
	if not os.path.isfile(file):
	    log.error("SRS file must exist in the current directory for source commands")
	    raise OSError  # XXX

	self.parseFile(file)

def _verifyAtHead(repos, headPkg, state):
    headVersion = repos.pkgLatestVersion(state.getName(),
					 state.getBranch())
    if not headVersion.equal(state.getVersion()):
	return False

    # make sure the files in this directory are based on the same
    # versions as those in the package at head
    bail = 0
    for (fileId, (path, version)) in state.iterFileList():
	if isinstance(version, versions.NewVersion):
	    assert(not headPkg.hasFile(fileId))
	    # new file, it shouldn't be in the old package at all
	else:
	    srcFileVersion = headPkg.getFile(fileId)[1]
	    if not version.equal(srcFileVersion):
		return False

    return True

def _getRecipeVersion(recipeFile):
    # load the recipe; we need this to figure out what version we're building
    try:
        loader = recipe.RecipeLoader(recipeFile)
    except recipe.RecipeFileError, e:
	log.error("unable to load recipe file %s: %s", recipeFile, str(e))
        return None

    if not loader:
	log.error("unable to load a valid recipe class from %s", recipeFile)
	return None

    assert(len(loader.values()) == 1)
    recipeClass = loader.values()[0]

    return recipeClass.version

def checkout(repos, cfg, dir, name, versionStr = None):
    # We have to be careful with branch nicknames.  First, we could get
    # multiple matches for a single package. Two, when a nickname uniquely
    # identifies a package we still need to make sure the state has the name of
    # the actual branch since empty branches yield objects whose versions are
    # on the parent branch.
    name += ":sources"
    try:
        trvList = helper.findPackage(repos, cfg.packagenamespace,
                                     cfg.installbranch, name,
                                     versionStr = versionStr)
    except helper.PackageNotFound, e:
        log.error(str(e))
        return
    if len(trvList) > 1:
	log.error("branch %s matches more then one version", versionStr)
	return
    trv = trvList[0]

    if not dir:
	dir = trv.getName().split(":")[-2]

    if not os.path.isdir(dir):
	try:
	    os.mkdir(dir)
	except OSError, err:
	    log.error("cannot create directory %s/%s: %s", os.getcwd(), dir,
                      str(err))
	    return

    branch = helper.fullBranchName(cfg.packagenamespace[1:], cfg.installbranch,
				   trv.getVersion(), versionStr)
    state = SourceState(trv.getName(), trv.getVersion(), branch)

    for (fileId, (path, version)) in trv.iterFileList():
	fullPath = dir + "/" + path
	(fileObj, contents) = repos.getFileVersion(fileId, version,
						   withContents = True)

	fileObj.restore(contents, fullPath, 1)

	state.addFile(fileId, path, version)

    state.write(dir + "/SRS")

def commit(repos):
    state = SourceStateFromFile("SRS")

    if isinstance(state.getVersion(), versions.NewVersion):
	# new package, so it shouldn't exist yet
	if repos.hasPackage(state.getName()):
	    log.error("%s is marked as a new package but it "
		      "already exists" % state.getName())
	    return
	srcPkg = None
    else:
	srcPkg = repos.getPackageVersion(state.getName(), state.getVersion())
	# update the version to one w/ a timestamp
	srcPkg.changeVersion(repos.pkgGetFullVersion(state.getName(),
			     srcPkg.getVersion()))

	if not _verifyAtHead(repos, srcPkg, state):
	    log.error("contents of working directory are not all "
		      "from the head of the branch; use update")
	    return

    recipeVersionStr = _getRecipeVersion(state.getRecipeFileName())
    if not recipeVersionStr: return

    if srcPkg:
	newVersion = helper.nextVersion(recipeVersionStr, srcPkg.getVersion(),
					state.getBranch(), binary = False)
    else:
	newVersion = helper.nextVersion(recipeVersionStr, None,
					state.getBranch(), binary = False)

    result = update.buildLocalChanges(repos, [(state, srcPkg, newVersion)])
    if not result: return

    (changeSet, ((isDifferent, newState),)) = result

    if not isDifferent:
	log.info("no changes have been made to commit")
    else:
	repos.commitChangeSet(changeSet)
	newState.write("SRS")

def diff(repos, versionStr = None):
    state = SourceStateFromFile("SRS")

    if versionStr:
	versionStr = state.expandVersionStr(versionStr)

	pkgList = helper.findPackage(repos, None, None, state.getName(),
				     versionStr)
	if len(pkgList) > 1:
	    log.error("%s specifies multiple versions" % versionStr)
	    return

	oldPackage = pkgList[0]
    else:
	oldPackage = repos.getPackageVersion(state.getName(),
					     state.getVersion())

    # update the version to one w/ a timestamp
    oldPackage.changeVersion(repos.pkgGetFullVersion(state.getName(),
			     oldPackage.getVersion()))

    result = update.buildLocalChanges(repos, [(state, oldPackage,
					       versions.NewVersion())])
    if not result: return

    (changeSet, ((isDifferent, newState),)) = result
    if not isDifferent: return

    packageChanges = changeSet.getNewPackageList()
    assert(len(packageChanges) == 1)
    pkgCs = packageChanges[0]

    for (fileId, path, newVersion) in pkgCs.getNewFileList():
	print "%s: new" % path

    for (fileId, path, newVersion) in pkgCs.getChangedFileList():
	if not path:
	    path = oldPackage.getFile(fileId)[0]
	    sys.stdout.write("%s" % path)
	else:
	    oldPath = oldPackage.getFile(fileId)[0]
	    sys.stdout.write("%s (aka %s)" % (path, oldPath))

	if not newVersion:
	    print
	    continue

	sys.stdout.write(": changed\n")

	csInfo = changeSet.getFileChange(fileId)
	print "    %s" % csInfo

	fileType = csInfo.split()[0]
	sha1 = csInfo.split()[1]
	if fileType == "f" and sha1 != "-":
	    (contType, contents) = changeSet.getFileContents(fileId)
	    if contType == changeset.ChangedFileTypes.diff:
		lines = contents.get().readlines()
		str = "    " + "    ".join(lines)
		print
		print str
		print

    for fileId in pkgCs.getOldFileList():
	path = oldPackage.getFile(fileId)[0]
	print "%s: removed" % path

def updateSrc(repos, versionStr = None):
    state = SourceStateFromFile("SRS")
    pkgName = state.getName()
    baseVersion = state.getVersion()

    if not versionStr:
	head = repos.getLatestPackage(pkgName, state.getBranch())
	newBranch = None
	headVersion = head.getVersion()
	if headVersion.equal(baseVersion):
	    log.info("working directory is already based on head of branch")
	    return
    else:
	versionStr = state.expandVersionStr(versionStr)

	pkgList = helper.findPackage(repos, None, None, pkgName, versionStr)
	if len(pkgList) > 1:
	    log.error("%s specifies multiple versions" % versionStr)
	    return

	head = pkgList[0]
	headVersion = head.getVersion()
	newBranch = helper.fullBranchName(None, None, headVersion, versionStr)

    changeSet = repos.createChangeSet([(pkgName, baseVersion, headVersion, 0)])

    packageChanges = changeSet.getNewPackageList()
    assert(len(packageChanges) == 1)
    pkgCs = packageChanges[0]

    fsJob = update.FilesystemJob(repos, changeSet,
				 { state.getName() : state }, "" )
    errList = fsJob.getErrorList()
    if errList:
	for err in errList: log.error(err)
    fsJob.apply()
    newPkgs = fsJob.getNewPackageList()
    assert(len(newPkgs) == 1)
    newState = newPkgs[0]

    if newState.getVersion().equal(pkgCs.getNewVersion()) and newBranch:
	newState.changeBranch(newBranch)

    newState.write("SRS")

def addFile(file):
    state = SourceStateFromFile("SRS")

    try:
	os.lstat(file)
    except OSError:
	log.error("files must be created before they can be added")
	return

    for (fileId, (path, version)) in state.iterFileList():
	if path == file:
	    log.error("file %s is already part of this source package" % path)
	    return

    fileId = cook.makeFileId(os.getcwd(), file)

    state.addFile(fileId, file, versions.NewVersion())
    state.write("SRS")

def removeFile(file):
    state = SourceStateFromFile("SRS")

    if not state.removeFilePath(file):
	log.error("file %s is not under management" % file)

    if os.path.exists(file):
	os.unlink(file)

    state.write("SRS")

def newPackage(repos, cfg, name):
    if name[0] != ":":
	name = cfg.packagenamespace + ":" + name
    name += ":sources"

    state = SourceState(name, versions.NewVersion(), cfg.defaultbranch)

    if repos and repos.hasPackage(name):
	log.error("package %s already exists" % name)
	return

    dir = name.split(":")[-2]
    if not os.path.isdir(dir):
	try:
	    os.mkdir(dir)
	except:
	    log.error("cannot create directory %s/%s", os.getcwd(), dir)
	    return

    state.write(dir + "/" + "SRS")

def renameFile(oldName, newName):
    state = SourceStateFromFile("SRS")

    if not os.path.exists(oldName):
	log.error("%s does not exist or is not a regular file" % oldName)
	return

    try:
	os.lstat(newName)
    except:
	pass
    else:
	log.error("%s already exists" % newName)
	return

    for (fileId, (path, version)) in state.iterFileList():
	if path == oldName:
	    os.rename(oldName, newName)
	    state.addFile(fileId, newName, version)
	    state.write("SRS")
	    return

    log.error("file %s is not under management" % oldName)
<|endoftext|>