<commit_before>"""Forms related to domains management."""

from django import forms
from django.core.urlresolvers import reverse
from django.http import QueryDict
from django.utils.translation import ugettext as _, ugettext_lazy

from modoboa.core.models import User
from modoboa.core import signals as core_signals
from modoboa.lib import events, parameters
from modoboa.lib.exceptions import ModoboaException, Conflict
from modoboa.lib.fields import DomainNameField
from modoboa.lib.form_utils import (
    YesNoField, WizardForm, WizardStep, DynamicForm, TabForms
)
from modoboa.lib.web_utils import render_to_json_response

from ..lib import check_if_domain_exists
from ..models import (
    Domain, DomainAlias, Mailbox, Alias
)

DOMAIN_TYPES = [
    ("domain", _("Domain")),
]

class DomainFormGeneral(forms.ModelForm, DynamicForm):

    """A form to create/edit a domain."""

    type = forms.ChoiceField(
        label=ugettext_lazy("Type"),
    )
    quota = forms.IntegerField(
        label=ugettext_lazy("Quota"),
        required=False,
        help_text=ugettext_lazy(
            "Default quota in MB applied to mailboxes. Leave empty to use the "
            "default value."
        )
    )
    aliases = DomainNameField(
        label=ugettext_lazy("Alias(es)"),
        required=False,
        help_text=ugettext_lazy(
            "Alias(es) of this domain. Indicate only one name per input, "
            "press ENTER to add a new input."
        )
    )

    class Meta:
        model = Domain
        fields = ("name", "type", "quota", "aliases", "enabled")

    def __init__(self, *args, **kwargs):
        self.oldname = None
        if "instance" in kwargs:
            self.oldname = kwargs["instance"].name
        super(DomainFormGeneral, self).__init__(*args, **kwargs)

        self.fields["type"].choices = (
            DOMAIN_TYPES + events.raiseQueryEvent("ExtraDomainTypes"))

        self.field_widths = {
            "quota": 3
        }

        if len(args) and isinstance(args[0], QueryDict):
            self._load_from_qdict(args[0], "aliases", DomainNameField)
        elif "instance" in kwargs:
            d = kwargs["instance"]
            for pos, dalias in enumerate(d.domainalias_set.all()):
                name = "aliases_%d" % (pos + 1)
                self._create_field(forms.CharField, name, dalias.name, 3)

    def clean_quota(self):
        if self.cleaned_data['quota'] is None:
            return int(parameters.get_admin('DEFAULT_DOMAIN_QUOTA'))
        if self.cleaned_data['quota'] < 0:
            raise forms.ValidationError(_("Must be a positive integer"))
        return self.cleaned_data['quota']

    def clean(self):
        """Custom fields validation.

        We want to prevent duplicate names between domains and domain
        aliases. Extensions have the possibility to declare other
        objects (see *CheckDomainName* event).

        The validation way is not very smart...
        """
        cleaned_data = super(DomainFormGeneral, self).clean()
        if self._errors:
            return cleaned_data
        name = cleaned_data["name"]
        label = check_if_domain_exists(
            name, [(DomainAlias, _('domain alias'))])
        if label is not None:
            self.add_error(
                "name", _("A %s with this name already exists")
                % unicode(label)
            )
        self.aliases = []
        for k in cleaned_data.keys():
            if not k.startswith("aliases"):
                continue
            if cleaned_data[k] == "":
                del cleaned_data[k]
                continue
            if cleaned_data[k] == name:
                self._errors[k] = self.error_class(
                    [_("A %s with this name already exists") % _("domain")]
                )
                del cleaned_data[k]
                continue
            label = check_if_domain_exists(
                cleaned_data[k], [(Domain, _("domain"))])
            if label is not None:
                self.add_error(
                    k, _("A %s with this name already exists")
                    % unicode(label)
                )
            else:
                self.aliases.append(cleaned_data[k])
        return cleaned_data

    def save(self, user, commit=True, domalias_post_create=False):
        """Custom save method.

        Updating a domain may have consequences on other objects
        (domain alias, mailbox, quota). The most tricky part concerns
        quotas update.

        """
        d = super(DomainFormGeneral, self).save(commit=False)
        if not commit:
            return d
        d.save()
        for dalias in d.domainalias_set.all():
            if dalias.name not in self.aliases:
                dalias.delete()
            else:
                self.aliases.remove(dalias.name)
        if self.aliases:
            core_signals.can_create_object.send(
                self.__class__, context=user, object_type="domain_aliases",
                count=len(self.aliases))
            core_signals.can_create_object.send(
                self.__class__, context=d, object_type="domain_aliases",
                count=len(self.aliases))
            for alias in self.aliases:
                if d.domainalias_set.filter(name=alias).exists():
                    continue
                options = {"creator": user} if domalias_post_create else {}
                DomainAlias(name=alias, target=d, enabled=d.enabled).save(
                    **options)
        return d

class DomainFormOptions(forms.Form):

    """A form containing options for domain creation."""

    create_dom_admin = YesNoField(
        label=ugettext_lazy("Create a domain administrator"),
        initial="no",
        help_text=ugettext_lazy(
            "Automatically create an administrator for this domain"
        )
    )

    dom_admin_username = forms.CharField(
        label=ugettext_lazy("Name"),
        initial="admin",
        help_text=ugettext_lazy(
            "The administrator's name. Don't include the domain's name here, "
            "it will be automatically appended."
        ),
        required=False
    )

    create_aliases = YesNoField(
        label=ugettext_lazy("Create aliases"),
        initial="yes",
        help_text=ugettext_lazy(
            "Automatically create standard aliases for this domain"
        ),
        required=False
    )

    def __init__(self, user, *args, **kwargs):
        super(DomainFormOptions, self).__init__(*args, **kwargs)
        result = events.raiseQueryEvent('UserCanSetRole', user, 'DomainAdmins')
        if False in result:
            self.fields = {}
            return
        if args:
            if args[0].get("create_dom_admin", "no") == "yes":
                self.fields["dom_admin_username"].required = True
                self.fields["create_aliases"].required = True

    def clean_dom_admin_username(self):
        """Ensure admin username is an email address."""
        if '@' in self.cleaned_data["dom_admin_username"]:
            raise forms.ValidationError(_("Invalid format"))
        return self.cleaned_data["dom_admin_username"]

    def save(self, *args, **kwargs):
        if not self.fields:
            return
        if self.cleaned_data["create_dom_admin"] == "no":
            return
        user = kwargs.pop("user")
        domain = kwargs.pop("domain")
        username = "%s@%s" % (
            self.cleaned_data["dom_admin_username"], domain.name)
        try:
            da = User.objects.get(username=username)
        except User.DoesNotExist:
            pass
        else:
            raise Conflict(_("User '%s' already exists") % username)
        core_signals.can_create_object.send(
            self.__class__, context=user, object_type="mailboxes")
        da = User(username=username, email=username, is_active=True)
        da.set_password("password")
        da.save()
        da.role = "DomainAdmins"
        da.post_create(user)

        mb = Mailbox(
            address=self.cleaned_data["dom_admin_username"], domain=domain,
            user=da, use_domain_quota=True
        )
        mb.set_quota(
            override_rules=user.has_perm("admin.change_domain"))
        mb.save(creator=user)

        if domain.type == "domain" and \
           self.cleaned_data["create_aliases"] == "yes":
            core_signals.can_create_object.send(
                self.__class__, context=user, object_type="mailbox_aliases")
            alias = Alias(
                address=u"postmaster@{}".format(domain.name),
                domain=domain, enabled=True
            )
            alias.save()
            alias.set_recipients([mb.full_address])
            alias.post_create(user)

        domain.add_admin(da)

class DomainForm(TabForms):

    """Domain edition form."""

    template_name = "admin/editdomainform.html"

    def __init__(self, request, *args, **kwargs):
        self.user = request.user
        self.forms = []
        self.domain = None
        if self.user.has_perm("admin.change_domain"):
            self.forms.append({
                "id": "general",
                "title": _("General"),
                "formtpl": "admin/domain_general_form.html",
                "cls": DomainFormGeneral,
                "mandatory": True
            })

        cbargs = [self.user]
        if "instances" in kwargs:
            self.domain = kwargs["instances"]["general"]
            cbargs += [self.domain]
        self.forms += events.raiseQueryEvent("ExtraDomainForm", *cbargs)
        if not self.forms:
            self.active_id = "admins"
        super(DomainForm, self).__init__(request, *args, **kwargs)

    def extra_context(self, context):
        domadmins = [u for u in self.domain.admins
                     if self.request.user.can_access(u) and not u.is_superuser]
        if not self.request.user.is_superuser:
            domadmins = [u for u in domadmins if u.group == "DomainAdmins"]
        context.update({
            "title": self.domain.name,
            "action": reverse(
                "admin:domain_change", args=[self.domain.pk]),
            "formid": "domform",
            "domain": self.domain,
            "domadmins": domadmins
        })

    def is_valid(self):
        """Custom validation.

        We just save the current name before it is potentially
        modified.

        """
        if "general" in self.instances:
            self.instances["general"].oldname = self.instances["general"].name
        return super(DomainForm, self).is_valid()

    def save(self):
        """Custom save method.

        As forms interact with each other, it is easier to make custom
        code to save them.
        """
        if not self.forms:
            return
        first_form = self.forms[0]["instance"]
        options = {}
        if isinstance(first_form, DomainFormGeneral):
            options["domalias_post_create"] = True
        first_form.save(self.request.user, **options)
        for f in self.forms[1:]:
            f["instance"].save(self.request.user)

    def done(self):
        if "general" in self.instances:
            events.raiseEvent("DomainModified", self.instances["general"])
        return render_to_json_response(_("Domain modified"))

class DomainWizard(WizardForm):

    """Domain creation wizard."""

    def __init__(self, request):
        super(DomainWizard, self).__init__(request)
        self.add_step(
            WizardStep(
                "general", DomainFormGeneral, _("General"),
                "admin/domain_general_form.html"
            )
        )
        steps = events.raiseQueryEvent("ExtraDomainWizardSteps")
        for step in steps:
            self.add_step(step)
        self.add_step(
            WizardStep(
                "options", DomainFormOptions, _("Options"),
                "admin/domain_options_form.html",
                [self.request.user]
            )
        )

    def extra_context(self, context):
        context.update({
            "title": _("New domain"),
            "action": reverse("admin:domain_add"),
            "formid": "domform"
        })

    def done(self):
        genform = self.first_step.form
        domain = genform.save(self.request.user)
        domain.post_create(self.request.user)
        for step in self.steps[1:]:
            if not step.check_access(self):
                continue
            try:
                step.form.save(user=self.request.user, domain=domain)
            except ModoboaException:
                from django.db import transaction
                transaction.rollback()
                raise
        return render_to_json_response(_("Domain created"))
<commit_msg>Remove useless code.<commit_after>"""Forms related to domains management."""

from django import forms
from django.core.urlresolvers import reverse
from django.http import QueryDict
from django.utils.translation import ugettext as _, ugettext_lazy

from modoboa.core.models import User
from modoboa.core import signals as core_signals
from modoboa.lib import events, parameters
from modoboa.lib.exceptions import ModoboaException, Conflict
from modoboa.lib.fields import DomainNameField
from modoboa.lib.form_utils import (
    YesNoField, WizardForm, WizardStep, DynamicForm, TabForms
)
from modoboa.lib.web_utils import render_to_json_response

from ..lib import check_if_domain_exists
from ..models import (
    Domain, DomainAlias, Mailbox, Alias
)

DOMAIN_TYPES = [
    ("domain", _("Domain")),
]

class DomainFormGeneral(forms.ModelForm, DynamicForm):

    """A form to create/edit a domain."""

    type = forms.ChoiceField(
        label=ugettext_lazy("Type"),
    )
    quota = forms.IntegerField(
        label=ugettext_lazy("Quota"),
        required=False,
        help_text=ugettext_lazy(
            "Default quota in MB applied to mailboxes. Leave empty to use the "
            "default value."
        )
    )
    aliases = DomainNameField(
        label=ugettext_lazy("Alias(es)"),
        required=False,
        help_text=ugettext_lazy(
            "Alias(es) of this domain. Indicate only one name per input, "
            "press ENTER to add a new input."
        )
    )

    class Meta:
        model = Domain
        fields = ("name", "type", "quota", "aliases", "enabled")

    def __init__(self, *args, **kwargs):
        self.oldname = None
        if "instance" in kwargs:
            self.oldname = kwargs["instance"].name
        super(DomainFormGeneral, self).__init__(*args, **kwargs)

        self.fields["type"].choices = (
            DOMAIN_TYPES + events.raiseQueryEvent("ExtraDomainTypes"))

        self.field_widths = {
            "quota": 3
        }

        if len(args) and isinstance(args[0], QueryDict):
            self._load_from_qdict(args[0], "aliases", DomainNameField)
        elif "instance" in kwargs:
            d = kwargs["instance"]
            for pos, dalias in enumerate(d.domainalias_set.all()):
                name = "aliases_%d" % (pos + 1)
                self._create_field(forms.CharField, name, dalias.name, 3)

    def clean_quota(self):
        if self.cleaned_data['quota'] is None:
            return int(parameters.get_admin('DEFAULT_DOMAIN_QUOTA'))
        if self.cleaned_data['quota'] < 0:
            raise forms.ValidationError(_("Must be a positive integer"))
        return self.cleaned_data['quota']

    def clean(self):
        """Custom fields validation.

        We want to prevent duplicate names between domains and domain
        aliases. Extensions have the possibility to declare other
        objects (see *CheckDomainName* event).

        The validation way is not very smart...
        """
        cleaned_data = super(DomainFormGeneral, self).clean()
        if self._errors:
            return cleaned_data
        name = cleaned_data["name"]
        label = check_if_domain_exists(
            name, [(DomainAlias, _('domain alias'))])
        if label is not None:
            self.add_error(
                "name", _("A %s with this name already exists")
                % unicode(label)
            )
        self.aliases = []
        for k in cleaned_data.keys():
            if not k.startswith("aliases"):
                continue
            if cleaned_data[k] == "":
                del cleaned_data[k]
                continue
            if cleaned_data[k] == name:
                self._errors[k] = self.error_class(
                    [_("A %s with this name already exists") % _("domain")]
                )
                del cleaned_data[k]
                continue
            label = check_if_domain_exists(
                cleaned_data[k], [(Domain, _("domain"))])
            if label is not None:
                self.add_error(
                    k, _("A %s with this name already exists")
                    % unicode(label)
                )
            else:
                self.aliases.append(cleaned_data[k])
        return cleaned_data

    def save(self, user, commit=True, domalias_post_create=False):
        """Custom save method.

        Updating a domain may have consequences on other objects
        (domain alias, mailbox, quota). The most tricky part concerns
        quotas update.

        """
        d = super(DomainFormGeneral, self).save(commit=False)
        if not commit:
            return d
        d.save()
        for dalias in d.domainalias_set.all():
            if dalias.name not in self.aliases:
                dalias.delete()
            else:
                self.aliases.remove(dalias.name)
        if self.aliases:
            core_signals.can_create_object.send(
                self.__class__, context=user, object_type="domain_aliases",
                count=len(self.aliases))
            core_signals.can_create_object.send(
                self.__class__, context=d, object_type="domain_aliases",
                count=len(self.aliases))
            for alias in self.aliases:
                if d.domainalias_set.filter(name=alias).exists():
                    continue
                options = {"creator": user} if domalias_post_create else {}
                DomainAlias(name=alias, target=d, enabled=d.enabled).save(
                    **options)
        return d

class DomainFormOptions(forms.Form):

    """A form containing options for domain creation."""

    create_dom_admin = YesNoField(
        label=ugettext_lazy("Create a domain administrator"),
        initial="no",
        help_text=ugettext_lazy(
            "Automatically create an administrator for this domain"
        )
    )

    dom_admin_username = forms.CharField(
        label=ugettext_lazy("Name"),
        initial="admin",
        help_text=ugettext_lazy(
            "The administrator's name. Don't include the domain's name here, "
            "it will be automatically appended."
        ),
        required=False
    )

    create_aliases = YesNoField(
        label=ugettext_lazy("Create aliases"),
        initial="yes",
        help_text=ugettext_lazy(
            "Automatically create standard aliases for this domain"
        ),
        required=False
    )

    def __init__(self, user, *args, **kwargs):
        super(DomainFormOptions, self).__init__(*args, **kwargs)
        result = events.raiseQueryEvent('UserCanSetRole', user, 'DomainAdmins')
        if False in result:
            self.fields = {}
            return
        if args:
            if args[0].get("create_dom_admin", "no") == "yes":
                self.fields["dom_admin_username"].required = True
                self.fields["create_aliases"].required = True

    def clean_dom_admin_username(self):
        """Ensure admin username is an email address."""
        if '@' in self.cleaned_data["dom_admin_username"]:
            raise forms.ValidationError(_("Invalid format"))
        return self.cleaned_data["dom_admin_username"]

    def save(self, *args, **kwargs):
        if not self.fields:
            return
        if self.cleaned_data["create_dom_admin"] == "no":
            return
        user = kwargs.pop("user")
        domain = kwargs.pop("domain")
        username = "%s@%s" % (
            self.cleaned_data["dom_admin_username"], domain.name)
        try:
            da = User.objects.get(username=username)
        except User.DoesNotExist:
            pass
        else:
            raise Conflict(_("User '%s' already exists") % username)
        core_signals.can_create_object.send(
            self.__class__, context=user, object_type="mailboxes")
        da = User(username=username, email=username, is_active=True)
        da.set_password("password")
        da.save()
        da.role = "DomainAdmins"
        da.post_create(user)

        mb = Mailbox(
            address=self.cleaned_data["dom_admin_username"], domain=domain,
            user=da, use_domain_quota=True
        )
        mb.set_quota(
            override_rules=user.has_perm("admin.change_domain"))
        mb.save(creator=user)

        if domain.type == "domain" and \
           self.cleaned_data["create_aliases"] == "yes":
            core_signals.can_create_object.send(
                self.__class__, context=user, object_type="mailbox_aliases")
            alias = Alias(
                address=u"postmaster@{}".format(domain.name),
                domain=domain, enabled=True
            )
            alias.save()
            alias.set_recipients([mb.full_address])
            alias.post_create(user)

        domain.add_admin(da)

class DomainForm(TabForms):

    """Domain edition form."""

    template_name = "admin/editdomainform.html"

    def __init__(self, request, *args, **kwargs):
        self.user = request.user
        self.forms = []
        self.domain = None
        if self.user.has_perm("admin.change_domain"):
            self.forms.append({
                "id": "general",
                "title": _("General"),
                "formtpl": "admin/domain_general_form.html",
                "cls": DomainFormGeneral,
                "mandatory": True
            })

        cbargs = [self.user]
        if "instances" in kwargs:
            self.domain = kwargs["instances"]["general"]
            cbargs += [self.domain]
        self.forms += events.raiseQueryEvent("ExtraDomainForm", *cbargs)
        if not self.forms:
            self.active_id = "admins"
        super(DomainForm, self).__init__(request, *args, **kwargs)

    def extra_context(self, context):
        domadmins = [u for u in self.domain.admins
                     if self.request.user.can_access(u) and not u.is_superuser]
        if not self.request.user.is_superuser:
            domadmins = [u for u in domadmins if u.group == "DomainAdmins"]
        context.update({
            "title": self.domain.name,
            "action": reverse(
                "admin:domain_change", args=[self.domain.pk]),
            "formid": "domform",
            "domain": self.domain,
            "domadmins": domadmins
        })

    def is_valid(self):
        """Custom validation.

        We just save the current name before it is potentially
        modified.

        """
        if "general" in self.instances:
            self.instances["general"].oldname = self.instances["general"].name
        return super(DomainForm, self).is_valid()

    def save(self):
        """Custom save method.

        As forms interact with each other, it is easier to make custom
        code to save them.
        """
        if not self.forms:
            return
        first_form = self.forms[0]["instance"]
        options = {}
        if isinstance(first_form, DomainFormGeneral):
            options["domalias_post_create"] = True
        first_form.save(self.request.user, **options)
        for f in self.forms[1:]:
            f["instance"].save(self.request.user)

    def done(self):
        if "general" in self.instances:
            events.raiseEvent("DomainModified", self.instances["general"])
        return render_to_json_response(_("Domain modified"))

class DomainWizard(WizardForm):

    """Domain creation wizard."""

    def __init__(self, request):
        super(DomainWizard, self).__init__(request)
        self.add_step(
            WizardStep(
                "general", DomainFormGeneral, _("General"),
                "admin/domain_general_form.html"
            )
        )
        steps = events.raiseQueryEvent("ExtraDomainWizardSteps")
        for step in steps:
            self.add_step(step)
        self.add_step(
            WizardStep(
                "options", DomainFormOptions, _("Options"),
                "admin/domain_options_form.html",
                [self.request.user]
            )
        )

    def extra_context(self, context):
        context.update({
            "title": _("New domain"),
            "action": reverse("admin:domain_add"),
            "formid": "domform"
        })

    def done(self):
        genform = self.first_step.form
        domain = genform.save(self.request.user)
        domain.post_create(self.request.user)
        for step in self.steps[1:]:
            if not step.check_access(self):
                continue
            step.form.save(user=self.request.user, domain=domain)
        return render_to_json_response(_("Domain created"))
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from argparse import ArgumentParser
import sys

from pypinyin import (__title__, __version__, pinyin, slug, lazy_pinyin,
                      load_single_dict, load_phrases_dict, NORMAL, TONE,
                      TONE2, INITIALS, FIRST_LETTER, FINALS, FINALS_TONE,
                      FINALS_TONE2)

# 支持命令行选项和参数
parser = ArgumentParser(description='convert chinese to pinyin.')
parser.add_argument('-V', '--version', action='version',
                    version='{0} {1}'.format(__title__, __version__))
# 要执行的函数名称
parser.add_argument('--func', help='function name (default: "pinyin")',
                    choices=['pinyin', 'slug'],
                    default='pinyin')
# 拼音风格
parser.add_argument('--style', help='pinyin style (default: "TONE")',
                    choices=['NORMAL', 'TONE', 'TONE2', 'INITIALS',
                             'FIRST_LETTER', 'FINALS', 'FINALS_TONE',
                             'FINALS_TONE2'], default='TONE')
parser.add_argument('--separator', help='slug separator (default: "-")',
                    default='-')
# 输出多音字
parser.add_argument('--heteronym', help='enable heteronym',
                    action='store_true')
# 要查询的汉字
parser.add_argument('hans', help='chinese string')

# python 3
try:
    unicode = unicode
except NameError:
    unicode = str

def main():
    # 获取命令行选项和参数
    options = parser.parse_args()
    hans = options.hans.decode(sys.stdin.encoding)
    func = globals()[options.func]
    style = globals()[options.style]
    heteronym = options.heteronym
    separator = options.separator

    kwargs = {
        'pinyin': {'heteronym': heteronym},
        'slug': {'heteronym': heteronym, 'separator': separator},
    }
    import logging
    # 禁用除 CRITICAL 外的日志消息
    logging.disable(logging.CRITICAL)
    result = func(hans, style=style, **kwargs[func.func_name])

    if result and isinstance(result, (list, tuple)):
        if isinstance(result[0], (list, tuple)):
            print(' '.join([','.join(s) for s in result]))
        else:
            print(result)
    else:
        print(result)

if __name__ == '__main__':
    main()
<commit_msg>修复命令行工具在 python 3 下的兼容问题<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from argparse import ArgumentParser
import sys

import pypinyin
from pypinyin import __title__, __version__

py3 = sys.version_info[0] == 3
if py3:
    unicode = str

def get_parser():
    parser = ArgumentParser(description='convert chinese to pinyin.')
    parser.add_argument('-V', '--version', action='version',
                        version='{0} {1}'.format(__title__, __version__))
    # 要执行的函数名称
    parser.add_argument('--func', help='function name (default: "pinyin")',
                        choices=['pinyin', 'slug'],
                        default='pinyin')
    # 拼音风格
    parser.add_argument('--style', help='pinyin style (default: "TONE")',
                        choices=['NORMAL', 'TONE', 'TONE2', 'INITIALS',
                                 'FIRST_LETTER', 'FINALS', 'FINALS_TONE',
                                 'FINALS_TONE2'], default='TONE')
    parser.add_argument('--separator', help='slug separator (default: "-")',
                        default='-')
    # 输出多音字
    parser.add_argument('--heteronym', help='enable heteronym',
                        action='store_true')
    # 要查询的汉字
    parser.add_argument('hans', help='chinese string')
    return parser

def main():
    import logging
    # 禁用除 CRITICAL 外的日志消息
    logging.disable(logging.CRITICAL)

    # 获取命令行选项和参数
    parser = get_parser()
    options = parser.parse_args()
    if py3:
        hans = options.hans
    else:
        hans = options.hans.decode(sys.stdin.encoding)
    func = getattr(pypinyin, options.func, 'pinyin')
    style = getattr(pypinyin, options.style, 'TONE')
    heteronym = options.heteronym
    separator = options.separator

    func_kwargs = {
        'pinyin': {'heteronym': heteronym},
        'slug': {'heteronym': heteronym, 'separator': separator},
    }
    if py3:
        kwargs = func_kwargs[func.__name__]
    else:
        kwargs = func_kwargs[func.func_name]
    result = func(hans, style=style, **kwargs)

    if result and isinstance(result, (list, tuple)):
        if isinstance(result[0], (list, tuple)):
            print(' '.join([','.join(s) for s in result]))
        else:
            print(result)
    else:
        print(result)

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


<commit_before>import salt.client
import re
import json
import salt.config
import xml.etree.ElementTree as ET
from salt.client.ssh.client import SSHClient
from jsonrpclib import Server

client = SSHClient()
local = salt.client.LocalClient()
master = salt.client.Caller()

__virtualname__ = 'nuts'

def __virtual__():
    return __virtualname__

def getCiscoXML(dst, param, cmd, attribut):
    value = master.cmd('cmd.run', "salt-ssh " + str(dst) + " -i -r '" + str(cmd) + " " + str(
        param) + " " + str(attribut) + " | format flash:nuts.odm' --roster-file=/etc/salt/roster")
    xml = value[value.index('<'):len(value)]
    return ET.fromstring(xml)

def returnMultiple(result):
    json_data = {}
    json_data["result"] = result
    json_data["resulttype"] = "multiple"
    return json.dumps(json_data)

def returnSingle(result):
    json_data = {}
    json_data["result"] = result
    json_data["resulttype"] = "single"
    return json.dumps(json_data)

def connectivity(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('ping -c 3 {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]*)% packet loss"
        r = re.compile(regex)
        m = r.search(text)
        return returnSingle((int(float(m.group(1))) < 100))
    elif os == "proxy":
        #at the moment there's a bug in the napalm-salt library which forces you to set all parameters fixed
        #https://github.com/saltstack/salt/pull/38577
        result = __salt__['net.ping'](param,'',255,2,100,10)
        #the absolute is needed because cisco returns -10 as packet_loss if they are not sent
        return returnSingle(result['result'] and abs(result['out']['success']['packet_loss']) != 10)

def traceroute(param):
    json_data = {}
    resultList = []
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('traceroute {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "[0-9]*  ([0-9\.]*) \("
        for m in re.finditer(regex, text):
            resultList.append(m.group(1))
        json_data["result"] = resultList
        json_data["resulttype"] = "multiple"
        return json.dumps(json_data)
    elif os == "proxy":
        result = __salt__['net.traceroute'](param)
        #TODO process resultlist
        if result['result']:
            return returnMultiple(result['out']['success'])

def bandwidth(dst, host, os, user, pwd):
    '''This function isn't working at the moment because it's not complient with the saltstack way'''
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    if os == "linux":
        local.cmd(dst, 'cmd.run', ['iperf3 -s -D -1'])
        result = local.cmd(host, 'cmd.run', ['iperf3 -c ' + dst])
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "Bytes  ([0-9\.]*) ([a-zA-Z]bits\/sec)([\s]*receiver)"
        r = re.compile(regex)
        m = r.search(text)
        return  returnSingle(float(m.group(1)) * 1000.0 * 1000.0)
    '''

def dnscheck(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('nslookup {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "(Name:[\s]*[a-z0-9.]*)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def dhcpcheck(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('dhcping -s {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "(Got answer)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def webresponse(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('curl -Is {} | head -n 1'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]{3} OK)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def portresponse(param,port):
    resultList = []
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('nmap -p {} {}'.format(port,param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]*)\/[a-z]* (open)"
        for m in re.finditer(regex, text):
            if m.group(2) == "open":
                resultList.append(m.group(1))
        return returnMultiple(resultList)

def checkuser():
    resultList = []
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['users.config']()
        if result['result']:
            return returnMultiple(result['out'].keys())

def checkversion():
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['net.facts']()
        if result['result']:
            if 'os_version' in result['out']:
                return returnSingle(result['out']['os_version'])

def checkospfneighbors(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_ip_ospf_neighbor}"
        tree = getCiscoXML(str(dst), "", "sh ip ospf neighbor", "")
        for id in tree.iter(tag=namespace + 'NeighborID'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def countospfneighbors(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    value = checkospfneighbors(dst, os, user, pwd)
    json_data = json.loads(value[value.index('{'):(value.index('}') + 1)])
    return returnSingle(len(json_data["result"]))
    '''

def checkospfneighborsstatus(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_ip_ospf_neighbor}"
        tree = getCiscoXML(dst, "", "sh ip ospf neighbor", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'NeighborID').text + ":" + id.find(namespace + 'State').text)
        return returnMultiple(resultList)
    '''

def stpinterfacestate(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Sts').text)
        return returnMultiple(resultList)
    '''

def stpinterfacerole(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Role').text)
        return returnMultiple(resultList)
    '''

def stpinterfacecost(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Cost').text)
        return returnMultiple(resultList)
    '''

def stprootid(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_root}"
        tree = getCiscoXML(dst, "", "sh spanning-tree root", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'RootID').text.split(' ')[1])
        return returnMultiple(resultList)
    '''

def stprootcost(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_root}"
        tree = getCiscoXML(dst, "", "sh spanning-tree root", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Cost').text)
        return returnMultiple(resultList)
    '''

def stpvlaninterfaces(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_vlan_*}"
        tree = getCiscoXML(dst, "", "sh spanning-tree vlan", "")
        for id in tree.iter(tag=namespace + 'Interface'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def stpvlanblockedports(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_blockedports}"
        tree = getCiscoXML(dst, "", "sh spanning-tree blockedports", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Name').text + ":" + id.find(namespace + 'BlockedInterfacesList').text)
        return returnMultiple(resultList)
    '''

def vlanports(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_vlan_id_*}"
        tree = getCiscoXML(dst, param, "sh vlan id", "")
        for id in tree.iter(tag=namespace + 'Ports'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def interfacestatus(dst, param, os, user, pwd):
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Status'):
            result = (id.text)
        return returnSingle(result)

def interfacevlan(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Vlan'):
            result = (id.text)
        return returnSingle(result)
    '''

def interfaceduplex(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Duplex'):
            result = id.text
        return returnSingle(result)
    '''

def interfacespeed(dst, param, os, user, pwd):
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Speed'):
            result = id.text
        return returnSingle(result)

def cdpneighbor(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_cdp_neighbors_*}"
        tree = getCiscoXML(dst, param, "sh cdp neighbors", "")
        for id in tree.iter(tag=namespace + 'DeviceID'):
            result = id.text
        return returnSingle(result)
    '''

def cdpneighborcount(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_cdp_neighbors}"
        tree = getCiscoXML(dst, "", "sh cdp neighbors", "")
        i = 0
        for id in tree.iter(tag=namespace + 'entry'):
            i += 1
        return returnSingle(i)
    '''

def arp(dst, param, os, user, pwd):
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_arp}"
        tree = getCiscoXML(dst, "", "sh arp", "")
        for id in tree.iter(tag=namespace + 'entry'):
            if id.find(namespace + 'Address').text == param:
                resultList.append(param + ":" + id.find(namespace + 'HardwareAddr').text)
                return returnSingle(resultList)
        resultList.append(param + ":-")
        return returnSingle(resultList)<commit_msg>Change interfacestatus, interfacespeed, arp to use salt-napalm<commit_after>import salt.client
import re
import json
import salt.config
import xml.etree.ElementTree as ET
from salt.client.ssh.client import SSHClient
from jsonrpclib import Server

client = SSHClient()
local = salt.client.LocalClient()
master = salt.client.Caller()

__virtualname__ = 'nuts'

def __virtual__():
    return __virtualname__

def getCiscoXML(dst, param, cmd, attribut):
    value = master.cmd('cmd.run', "salt-ssh " + str(dst) + " -i -r '" + str(cmd) + " " + str(
        param) + " " + str(attribut) + " | format flash:nuts.odm' --roster-file=/etc/salt/roster")
    xml = value[value.index('<'):len(value)]
    return ET.fromstring(xml)

def returnMultiple(result):
    json_data = {}
    json_data["result"] = result
    json_data["resulttype"] = "multiple"
    return json.dumps(json_data)

def returnSingle(result):
    json_data = {}
    json_data["result"] = result
    json_data["resulttype"] = "single"
    return json.dumps(json_data)

def connectivity(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('ping -c 3 {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]*)% packet loss"
        r = re.compile(regex)
        m = r.search(text)
        return returnSingle((int(float(m.group(1))) < 100))
    elif os == "proxy":
        #at the moment there's a bug in the napalm-salt library which forces you to set all parameters fixed
        #https://github.com/saltstack/salt/pull/38577
        result = __salt__['net.ping'](param,'',255,2,100,10)
        #the absolute is needed because cisco returns -10 as packet_loss if they are not sent
        return returnSingle(result['result'] and abs(result['out']['success']['packet_loss']) != 10)

def traceroute(param):
    json_data = {}
    resultList = []
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('traceroute {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "[0-9]*  ([0-9\.]*) \("
        for m in re.finditer(regex, text):
            resultList.append(m.group(1))
        json_data["result"] = resultList
        json_data["resulttype"] = "multiple"
        return json.dumps(json_data)
    elif os == "proxy":
        result = __salt__['net.traceroute'](param)
        #TODO process resultlist
        if result['result']:
            return returnMultiple(result['out']['success'])

def bandwidth(dst, host, os, user, pwd):
    '''This function isn't working at the moment because it's not complient with the saltstack way'''
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    if os == "linux":
        local.cmd(dst, 'cmd.run', ['iperf3 -s -D -1'])
        result = local.cmd(host, 'cmd.run', ['iperf3 -c ' + dst])
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "Bytes  ([0-9\.]*) ([a-zA-Z]bits\/sec)([\s]*receiver)"
        r = re.compile(regex)
        m = r.search(text)
        return  returnSingle(float(m.group(1)) * 1000.0 * 1000.0)
    '''

def dnscheck(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('nslookup {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "(Name:[\s]*[a-z0-9.]*)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def dhcpcheck(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('dhcping -s {}'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "(Got answer)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def webresponse(param):
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('curl -Is {} | head -n 1'.format(param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]{3} OK)"
        r = re.compile(regex)
        return returnSingle(bool(re.search(regex, text)))

def portresponse(param,port):
    resultList = []
    os = __grains__['os_family']
    if os == "Debian":
        result = __salt__['cmd.run']('nmap -p {} {}'.format(port,param))
        text = bytes(result).decode(encoding="utf-8", errors='ignore')
        regex = "([0-9]*)\/[a-z]* (open)"
        for m in re.finditer(regex, text):
            if m.group(2) == "open":
                resultList.append(m.group(1))
        return returnMultiple(resultList)

def checkuser():
    resultList = []
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['users.config']()
        if result['result']:
            return returnMultiple(result['out'].keys())

def checkversion():
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['net.facts']()
        if result['result']:
            if 'os_version' in result['out']:
                return returnSingle(result['out']['os_version'])

def checkospfneighbors(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_ip_ospf_neighbor}"
        tree = getCiscoXML(str(dst), "", "sh ip ospf neighbor", "")
        for id in tree.iter(tag=namespace + 'NeighborID'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def countospfneighbors(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    value = checkospfneighbors(dst, os, user, pwd)
    json_data = json.loads(value[value.index('{'):(value.index('}') + 1)])
    return returnSingle(len(json_data["result"]))
    '''

def checkospfneighborsstatus(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_ip_ospf_neighbor}"
        tree = getCiscoXML(dst, "", "sh ip ospf neighbor", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'NeighborID').text + ":" + id.find(namespace + 'State').text)
        return returnMultiple(resultList)
    '''

def stpinterfacestate(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Sts').text)
        return returnMultiple(resultList)
    '''

def stpinterfacerole(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Role').text)
        return returnMultiple(resultList)
    '''

def stpinterfacecost(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_interface_*}"
        tree = getCiscoXML(dst, param, "sh spanning-tree interface", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Cost').text)
        return returnMultiple(resultList)
    '''

def stprootid(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_root}"
        tree = getCiscoXML(dst, "", "sh spanning-tree root", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'RootID').text.split(' ')[1])
        return returnMultiple(resultList)
    '''

def stprootcost(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_root}"
        tree = getCiscoXML(dst, "", "sh spanning-tree root", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Vlan').text + ":" + id.find(namespace + 'Cost').text)
        return returnMultiple(resultList)
    '''

def stpvlaninterfaces(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_vlan_*}"
        tree = getCiscoXML(dst, "", "sh spanning-tree vlan", "")
        for id in tree.iter(tag=namespace + 'Interface'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def stpvlanblockedports(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_spanning-tree_blockedports}"
        tree = getCiscoXML(dst, "", "sh spanning-tree blockedports", "")
        for id in tree.iter(tag=namespace + 'entry'):
            resultList.append(
                id.find(namespace + 'Name').text + ":" + id.find(namespace + 'BlockedInterfacesList').text)
        return returnMultiple(resultList)
    '''

def vlanports(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    resultList = []
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_vlan_id_*}"
        tree = getCiscoXML(dst, param, "sh vlan id", "")
        for id in tree.iter(tag=namespace + 'Ports'):
            resultList.append(id.text)
        return returnMultiple(resultList)
    '''

def interfacestatus(param):
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['net.interfaces']()
        if result['result']:
            if param in result['out']:
                return returnSingle(result['out'][param]['is_up'])

def interfacevlan(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Vlan'):
            result = (id.text)
        return returnSingle(result)
    '''

def interfaceduplex(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_interface_*_status}"
        tree = getCiscoXML(dst, param, "sh interface", " status")
        for id in tree.iter(tag=namespace + 'Duplex'):
            result = id.text
        return returnSingle(result)
    '''

def interfacespeed(param):
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['net.interfaces']()
        if result['result']:
            if param in result['out']:
                return returnSingle(result['out'][param]['speed'])

def cdpneighbor(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    result = ""
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_cdp_neighbors_*}"
        tree = getCiscoXML(dst, param, "sh cdp neighbors", "")
        for id in tree.iter(tag=namespace + 'DeviceID'):
            result = id.text
        return returnSingle(result)
    '''

def cdpneighborcount(dst, param, os, user, pwd):
    raise NotImplementedError('This function isn\'t implemented right now')
    '''
    if os == 'ios':
        namespace = "{ODM://flash:nuts.odm//show_cdp_neighbors}"
        tree = getCiscoXML(dst, "", "sh cdp neighbors", "")
        i = 0
        for id in tree.iter(tag=namespace + 'entry'):
            i += 1
        return returnSingle(i)
    '''

def arp(param):
    os = __grains__['os_family']
    if os == "proxy":
        result = __salt__['net.arp']('',param)
        if result['result']:
            for entry in result['out']:
                if(entry['ip'] == param):
                    return returnSingle(entry['mac'])<|endoftext|>


<commit_before>#===============================================================================
# Copyright (C) 2011-2012 by Andrew Moffat
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#===============================================================================

import subprocess as subp
import inspect
import sys
import traceback
import os
import re
from glob import glob
import shlex
import warnings

VERSION = "0.5"
PROJECT_URL = "https://github.com/amoffat/pbs"
IS_PY3 = sys.version_info[0] == 3

if IS_PY3: raw_input = input

class ErrorReturnCode(Exception):
    truncate_cap = 200

    def __init__(self, full_cmd, stdout, stderr):
        self.full_cmd = full_cmd
        self.stdout = stdout
        self.stderr = stderr

        tstdout = self.stdout[:self.truncate_cap]
        out_delta = len(self.stdout) - len(tstdout)
        if out_delta: tstdout += "... (%d more, please see e.stdout)" % out_delta

        tstderr = self.stderr[:self.truncate_cap]
        err_delta = len(self.stderr) - len(tstderr)
        if err_delta: tstderr += "... (%d more, please see e.stderr)" % err_delta

        msg = "\n\nRan: %r\n\nSTDOUT:\n\n  %s\nSTDERR:\n\n  %s" %\
            (full_cmd, tstdout, tstderr)
        super(ErrorReturnCode, self).__init__(msg)

class CommandNotFound(Exception): pass

rc_exc_regex = re.compile("ErrorReturnCode_(\d+)")
rc_exc_cache = {}

def get_rc_exc(rc):
    try: return rc_exc_cache[rc]
    except KeyError: pass

    name = "ErrorReturnCode_%d" % rc
    exc = type(name, (ErrorReturnCode,), {})
    rc_exc_cache[name] = exc
    return exc

def which(program):
    def is_exe(fpath):
        return os.path.exists(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program): return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def resolve_program(program):
    path = which(program)
    if not path:
        # our actual command might have a dash in it, but we can't call
        # that from python (we have to use underscores), so we'll check
        # if a dash version of our underscore command exists and use that
        # if it does
        if "_" in program: path = which(program.replace("_", "-"))
        if not path: return None
    return path

class Command(object):
    prepend_stack = []

    @classmethod
    def create(cls, program, raise_exc=True):
        path = resolve_program(program)
        if not path:
            if raise_exc: raise CommandNotFound(program)
            else: return None
        return cls(path)

    def __init__(self, path):
        self.path = path

        self.process = None
        self._stdout = None
        self._stderr = None

        self.call_args = {
            "bg": False, # run command in background
            "with": False, # prepend the command to every command after it
        }

    def __getattr__(self, p):
        return getattr(self.stdout, p)

    @property
    def stdout(self):
        if self.call_args["bg"]: self.wait()
        return self._stdout

    @property
    def stderr(self):
        if self.call_args["bg"]: self.wait()
        return self._stderr

    def wait(self):
        if self.process.returncode is not None: return
        self._stdout, self._stderr = self.process.communicate()
        rc = self.process.wait()

        if rc != 0: raise get_rc_exc(rc)(self.stdout, self.stderr)
        return self

    def __repr__(self):
        return str(self)

    def __long__(self):
        return long(str(self).strip())

    def __float__(self):
        return float(str(self).strip())

    def __int__(self):
        return int(str(self).strip())

    def __str__(self):
        if IS_PY3: return self.__unicode__()
        else: return unicode(self).encode('utf-8')

    def __unicode__(self):
        if self.process: return self.stdout.decode('utf-8') # byte string
        else: return self.path

    def __enter__(self):
        if not self.call_args["with"]:
            Command.prepend_stack.append([self.path])

    def __exit__(self, typ, value, traceback):
        if Command.prepend_stack:
            Command.prepend_stack.pop()

    def __call__(self, *args, **kwargs):
        kwargs = kwargs.copy()
        args = list(args)
        stdin = None
        processed_args = []
        cmd = []

        # aggregate any with contexts
        for prepend in self.prepend_stack:
            cmd.extend(prepend)

        cmd.append(self.path)

        # pull out the pbs-specific arguments (arguments that are not to be
        # passed to the commands
        for parg, default in self.call_args.items():
            key = "_" + parg
            self.call_args[parg] = default
            if key in kwargs:
                self.call_args[parg] = kwargs[key]
                del kwargs[key]

        # check if we're piping via composition
        stdin = subp.PIPE
        actual_stdin = None
        if args:
            first_arg = args.pop(0)
            if isinstance(first_arg, Command):
                # it makes sense that if the input pipe of a command is running
                # in the background, then this command should run in the
                # background as well
                if first_arg.call_args["bg"]:
                    self.call_args["bg"] = True
                    stdin = first_arg.process.stdout
                else:
                    actual_stdin = first_arg.stdout
            else: args.insert(0, first_arg)

        # aggregate the position arguments
        for arg in args:
            # i should've commented why i did this, because now i don't
            # remember.  for some reason, some command was more natural
            # taking a list?
            if isinstance(arg, (list, tuple)):
                for sub_arg in arg: processed_args.append(str(sub_arg))
            else: processed_args.append(str(arg))

        # aggregate the keyword arguments
        for k,v in kwargs.items():
            # we're passing a short arg as a kwarg, example:
            # cut(d="\t")
            if len(k) == 1:
                if v is True: arg = "-"+k
                else: arg = "-%s %r" % (k, v)

            # we're doing a long arg
            else:
                k = k.replace("_", "-")

                if v is True: arg = "--"+k
                else: arg = "--%s=%s" % (k, v)
            processed_args.append(arg)

        # makes sure our arguments are broken up correctly
        split_args = shlex.split(" ".join(processed_args))

        # now glob-expand each arg and compose the final list
        final_args = []
        for arg in split_args:
            expanded = glob(arg)
            if expanded: final_args.extend(expanded)
            else: final_args.append(arg)

        cmd.extend(final_args)
        # for debugging
        self._command_ran = " ".join(cmd)

        # with contexts shouldn't run at all yet, they prepend
        # to every command in the context
        if self.call_args["with"]:
            Command.prepend_stack.append(cmd)
            return self

        self.process = subp.Popen(cmd, shell=False, env=os.environ,
            stdin=stdin, stdout=subp.PIPE, stderr=subp.PIPE)

        if self.call_args["bg"]: return self

        self._stdout, self._stderr = self.process.communicate(actual_stdin)
        rc = self.process.wait()

        if rc != 0: raise get_rc_exc(rc)(self._command_ran, self.stdout, self.stderr)
        return self

# this class is used directly when we do a "from pbs import *".  it allows
# lookups to names that aren't found in the global scope to be searched
# for as a program.  for example, if "ls" isn't found in the program's
# scope, we consider it a system program and try to find it.
class Environment(dict):
    def __init__(self, *args, **kwargs):
        dict.__init__(self, *args, **kwargs)

        self["Command"] = Command
        self["CommandNotFound"] = CommandNotFound
        self["ErrorReturnCode"] = ErrorReturnCode
        self["ARGV"] = sys.argv[1:]
        for i, arg in enumerate(sys.argv):
            self["ARG%d" % i] = arg

        # this needs to be last
        self["env"] = os.environ

    def __setitem__(self, k, v):
        # are we altering an environment variable?
        if "env" in self and k in self["env"]: self["env"][k] = v
        # no?  just setting a regular name
        else: dict.__setitem__(self, k, v)

    def __missing__(self, k):
        # the only way we'd get to here is if we've tried to
        # import * from a repl.  so, raise an exception, since
        # that's really the only sensible thing to do
        if k == "__all__":
            raise RuntimeError("Cannot import * from the commandline, please \
see \"Limitations\" here: %s" % PROJECT_URL)

        # if we end with "_" just go ahead and skip searching
        # our namespace for python stuff.  this was mainly for the
        # command "id", which is a popular program for finding
        # if a user exists, but also a python function for getting
        # the address of an object.  so can call the python
        # version by "id" and the program version with "id_"
        if not k.endswith("_"):
            # check if we're naming a dynamically generated ReturnCode exception
            try: return rc_exc_cache[k]
            except KeyError:
                m = rc_exc_regex.match(k)
                if m: return get_rc_exc(int(m.group(1)))

            # are we naming a commandline argument?
            if k.startswith("ARG"):
                return None

            # is it a builtin?
            try: return getattr(self["__builtins__"], k)
            except AttributeError: pass
        elif not k.startswith("_"): k = k.rstrip("_")

        # how about an environment variable?
        try: return os.environ[k]
        except KeyError: pass

        # is it a custom builtin?
        builtin = getattr(self, "b_"+k, None)
        if builtin: return builtin

        # it must be a command then
        return Command.create(k)

    def b_echo(self, *args, **kwargs):
        out = Command("echo")(*args, **kwargs)
        print(out)
        return out

    def b_cd(self, path):
        os.chdir(path)

    def b_which(self, program):
        return which(program)

def run_repl(env):
    banner = "\n>> PBS v{version}\n>> https://github.com/amoffat/pbs\n"

    print(banner.format(version=VERSION))
    while True:
        try: line = raw_input("pbs> ")
        except (ValueError, EOFError): break

        try: exec(compile(line, "<dummy>", "single"), env, env)
        except SystemExit: break
        except: print(traceback.format_exc())

    # cleans up our last line
    print('')

# this is a thin wrapper around THIS module (we patch sys.modules[__name__]).
# this is in the case that the user does a "from pbs import whatever"
# in other words, they only want to import certain programs, not the whole
# system PATH worth of commands.  in this case, we just proxy the
# import lookup to our Environment class
class SelfWrapper(object):
    def __init__(self, self_module):
        self.self_module = self_module
        self.env = Environment(globals())

    def __getattr__(self, name):
        return self.env[name]

# we're being run as a stand-alone script, fire up a REPL
if __name__ == "__main__":
    globs = globals()
    f_globals = {}
    for k in ["__builtins__", "__doc__", "__name__", "__package__"]:
        f_globals[k] = globs[k]
    env = Environment(f_globals)
    run_repl(env)

# we're being imported from somewhere
else:
    frame, script, line, module, code, index = inspect.stack()[1]
    env = Environment(frame.f_globals)

    # are we being imported from a REPL? don't allow
    if script == "<stdin>" or script == "<ipython console>":
        self = sys.modules[__name__]
        sys.modules[__name__] = SelfWrapper(self)

    # we're being imported from a script
    else:

        # we need to analyze how we were imported
        with open(script, "r") as h: source = h.readlines()
        import_line = source[line-1]

        # this it the most magical choice.  basically we're trying to import
        # all of the system programs into our script.  the only way to do
        # this is going to be to exec the source in modified global scope.
        # there might be a less magical way to do this...
        if "*" in import_line:
            # do not let us import * from anywhere but a stand-alone script
            if frame.f_globals["__name__"] != "__main__":
                raise RuntimeError("Cannot import * from anywhere other than \
a stand-alone script.  Do a 'from pbs import program' instead. Please see \
\"Limitations\" here: %s" % PROJECT_URL)

            warnings.warn("Importing * from pbs is magical and therefore has \
some limitations.  Please become familiar with them under \"Limitations\" \
here: %s  To avoid this warning, use a warning filter or import your \
programs directly with \"from pbs import <program>\"" % PROJECT_URL,
RuntimeWarning, stacklevel=2)

            # we avoid recursion by removing the line that imports us :)
            source = "".join(source[line:])

            exit_code = 0
            try: exec(source, env, env)
            except SystemExit as e: exit_code = e.code
            except: print(traceback.format_exc())

            # we exit so we don't actually run the script that we were imported
            # from (which would be running it "again", since we just executed
            # the script with exec
            exit(exit_code)

        # this is the least magical choice.  we're importing either a
        # selection of programs or we're just importing the pbs module.
        # in this case, let's just wrap ourselves with a module that has
        # __getattr__ so our program lookups can be done there
        else:
            self = sys.modules[__name__]
            sys.modules[__name__] = SelfWrapper(self)
<commit_msg>making the check for a python shell (regular or ipython) a little more general<commit_after>#===============================================================================
# Copyright (C) 2011-2012 by Andrew Moffat
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#===============================================================================

import subprocess as subp
import inspect
import sys
import traceback
import os
import re
from glob import glob
import shlex
import warnings

VERSION = "0.5"
PROJECT_URL = "https://github.com/amoffat/pbs"
IS_PY3 = sys.version_info[0] == 3

if IS_PY3: raw_input = input

class ErrorReturnCode(Exception):
    truncate_cap = 200

    def __init__(self, full_cmd, stdout, stderr):
        self.full_cmd = full_cmd
        self.stdout = stdout
        self.stderr = stderr

        tstdout = self.stdout[:self.truncate_cap]
        out_delta = len(self.stdout) - len(tstdout)
        if out_delta: tstdout += "... (%d more, please see e.stdout)" % out_delta

        tstderr = self.stderr[:self.truncate_cap]
        err_delta = len(self.stderr) - len(tstderr)
        if err_delta: tstderr += "... (%d more, please see e.stderr)" % err_delta

        msg = "\n\nRan: %r\n\nSTDOUT:\n\n  %s\nSTDERR:\n\n  %s" %\
            (full_cmd, tstdout, tstderr)
        super(ErrorReturnCode, self).__init__(msg)

class CommandNotFound(Exception): pass

rc_exc_regex = re.compile("ErrorReturnCode_(\d+)")
rc_exc_cache = {}

def get_rc_exc(rc):
    try: return rc_exc_cache[rc]
    except KeyError: pass

    name = "ErrorReturnCode_%d" % rc
    exc = type(name, (ErrorReturnCode,), {})
    rc_exc_cache[name] = exc
    return exc

def which(program):
    def is_exe(fpath):
        return os.path.exists(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program): return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def resolve_program(program):
    path = which(program)
    if not path:
        # our actual command might have a dash in it, but we can't call
        # that from python (we have to use underscores), so we'll check
        # if a dash version of our underscore command exists and use that
        # if it does
        if "_" in program: path = which(program.replace("_", "-"))
        if not path: return None
    return path

class Command(object):
    prepend_stack = []

    @classmethod
    def create(cls, program, raise_exc=True):
        path = resolve_program(program)
        if not path:
            if raise_exc: raise CommandNotFound(program)
            else: return None
        return cls(path)

    def __init__(self, path):
        self.path = path

        self.process = None
        self._stdout = None
        self._stderr = None

        self.call_args = {
            "bg": False, # run command in background
            "with": False, # prepend the command to every command after it
        }

    def __getattr__(self, p):
        return getattr(self.stdout, p)

    @property
    def stdout(self):
        if self.call_args["bg"]: self.wait()
        return self._stdout

    @property
    def stderr(self):
        if self.call_args["bg"]: self.wait()
        return self._stderr

    def wait(self):
        if self.process.returncode is not None: return
        self._stdout, self._stderr = self.process.communicate()
        rc = self.process.wait()

        if rc != 0: raise get_rc_exc(rc)(self.stdout, self.stderr)
        return self

    def __repr__(self):
        return str(self)

    def __long__(self):
        return long(str(self).strip())

    def __float__(self):
        return float(str(self).strip())

    def __int__(self):
        return int(str(self).strip())

    def __str__(self):
        if IS_PY3: return self.__unicode__()
        else: return unicode(self).encode('utf-8')

    def __unicode__(self):
        if self.process: return self.stdout.decode('utf-8') # byte string
        else: return self.path

    def __enter__(self):
        if not self.call_args["with"]:
            Command.prepend_stack.append([self.path])

    def __exit__(self, typ, value, traceback):
        if Command.prepend_stack:
            Command.prepend_stack.pop()

    def __call__(self, *args, **kwargs):
        kwargs = kwargs.copy()
        args = list(args)
        stdin = None
        processed_args = []
        cmd = []

        # aggregate any with contexts
        for prepend in self.prepend_stack:
            cmd.extend(prepend)

        cmd.append(self.path)

        # pull out the pbs-specific arguments (arguments that are not to be
        # passed to the commands
        for parg, default in self.call_args.items():
            key = "_" + parg
            self.call_args[parg] = default
            if key in kwargs:
                self.call_args[parg] = kwargs[key]
                del kwargs[key]

        # check if we're piping via composition
        stdin = subp.PIPE
        actual_stdin = None
        if args:
            first_arg = args.pop(0)
            if isinstance(first_arg, Command):
                # it makes sense that if the input pipe of a command is running
                # in the background, then this command should run in the
                # background as well
                if first_arg.call_args["bg"]:
                    self.call_args["bg"] = True
                    stdin = first_arg.process.stdout
                else:
                    actual_stdin = first_arg.stdout
            else: args.insert(0, first_arg)

        # aggregate the position arguments
        for arg in args:
            # i should've commented why i did this, because now i don't
            # remember.  for some reason, some command was more natural
            # taking a list?
            if isinstance(arg, (list, tuple)):
                for sub_arg in arg: processed_args.append(str(sub_arg))
            else: processed_args.append(str(arg))

        # aggregate the keyword arguments
        for k,v in kwargs.items():
            # we're passing a short arg as a kwarg, example:
            # cut(d="\t")
            if len(k) == 1:
                if v is True: arg = "-"+k
                else: arg = "-%s %r" % (k, v)

            # we're doing a long arg
            else:
                k = k.replace("_", "-")

                if v is True: arg = "--"+k
                else: arg = "--%s=%s" % (k, v)
            processed_args.append(arg)

        # makes sure our arguments are broken up correctly
        split_args = shlex.split(" ".join(processed_args))

        # now glob-expand each arg and compose the final list
        final_args = []
        for arg in split_args:
            expanded = glob(arg)
            if expanded: final_args.extend(expanded)
            else: final_args.append(arg)

        cmd.extend(final_args)
        # for debugging
        self._command_ran = " ".join(cmd)

        # with contexts shouldn't run at all yet, they prepend
        # to every command in the context
        if self.call_args["with"]:
            Command.prepend_stack.append(cmd)
            return self

        self.process = subp.Popen(cmd, shell=False, env=os.environ,
            stdin=stdin, stdout=subp.PIPE, stderr=subp.PIPE)

        if self.call_args["bg"]: return self

        self._stdout, self._stderr = self.process.communicate(actual_stdin)
        rc = self.process.wait()

        if rc != 0: raise get_rc_exc(rc)(self._command_ran, self.stdout, self.stderr)
        return self

# this class is used directly when we do a "from pbs import *".  it allows
# lookups to names that aren't found in the global scope to be searched
# for as a program.  for example, if "ls" isn't found in the program's
# scope, we consider it a system program and try to find it.
class Environment(dict):
    def __init__(self, *args, **kwargs):
        dict.__init__(self, *args, **kwargs)

        self["Command"] = Command
        self["CommandNotFound"] = CommandNotFound
        self["ErrorReturnCode"] = ErrorReturnCode
        self["ARGV"] = sys.argv[1:]
        for i, arg in enumerate(sys.argv):
            self["ARG%d" % i] = arg

        # this needs to be last
        self["env"] = os.environ

    def __setitem__(self, k, v):
        # are we altering an environment variable?
        if "env" in self and k in self["env"]: self["env"][k] = v
        # no?  just setting a regular name
        else: dict.__setitem__(self, k, v)

    def __missing__(self, k):
        # the only way we'd get to here is if we've tried to
        # import * from a repl.  so, raise an exception, since
        # that's really the only sensible thing to do
        if k == "__all__":
            raise RuntimeError("Cannot import * from the commandline, please \
see \"Limitations\" here: %s" % PROJECT_URL)

        # if we end with "_" just go ahead and skip searching
        # our namespace for python stuff.  this was mainly for the
        # command "id", which is a popular program for finding
        # if a user exists, but also a python function for getting
        # the address of an object.  so can call the python
        # version by "id" and the program version with "id_"
        if not k.endswith("_"):
            # check if we're naming a dynamically generated ReturnCode exception
            try: return rc_exc_cache[k]
            except KeyError:
                m = rc_exc_regex.match(k)
                if m: return get_rc_exc(int(m.group(1)))

            # are we naming a commandline argument?
            if k.startswith("ARG"):
                return None

            # is it a builtin?
            try: return getattr(self["__builtins__"], k)
            except AttributeError: pass
        elif not k.startswith("_"): k = k.rstrip("_")

        # how about an environment variable?
        try: return os.environ[k]
        except KeyError: pass

        # is it a custom builtin?
        builtin = getattr(self, "b_"+k, None)
        if builtin: return builtin

        # it must be a command then
        return Command.create(k)

    def b_echo(self, *args, **kwargs):
        out = Command("echo")(*args, **kwargs)
        print(out)
        return out

    def b_cd(self, path):
        os.chdir(path)

    def b_which(self, program):
        return which(program)

def run_repl(env):
    banner = "\n>> PBS v{version}\n>> https://github.com/amoffat/pbs\n"

    print(banner.format(version=VERSION))
    while True:
        try: line = raw_input("pbs> ")
        except (ValueError, EOFError): break

        try: exec(compile(line, "<dummy>", "single"), env, env)
        except SystemExit: break
        except: print(traceback.format_exc())

    # cleans up our last line
    print('')

# this is a thin wrapper around THIS module (we patch sys.modules[__name__]).
# this is in the case that the user does a "from pbs import whatever"
# in other words, they only want to import certain programs, not the whole
# system PATH worth of commands.  in this case, we just proxy the
# import lookup to our Environment class
class SelfWrapper(object):
    def __init__(self, self_module):
        self.self_module = self_module
        self.env = Environment(globals())

    def __getattr__(self, name):
        return self.env[name]

# we're being run as a stand-alone script, fire up a REPL
if __name__ == "__main__":
    globs = globals()
    f_globals = {}
    for k in ["__builtins__", "__doc__", "__name__", "__package__"]:
        f_globals[k] = globs[k]
    env = Environment(f_globals)
    run_repl(env)

# we're being imported from somewhere
else:
    frame, script, line, module, code, index = inspect.stack()[1]
    env = Environment(frame.f_globals)

    # are we being imported from a REPL?
    if script.startswith("<") and script.endswith(">") :
        self = sys.modules[__name__]
        sys.modules[__name__] = SelfWrapper(self)

    # we're being imported from a script
    else:

        # we need to analyze how we were imported
        with open(script, "r") as h: source = h.readlines()
        import_line = source[line-1]

        # this it the most magical choice.  basically we're trying to import
        # all of the system programs into our script.  the only way to do
        # this is going to be to exec the source in modified global scope.
        # there might be a less magical way to do this...
        if "*" in import_line:
            # do not let us import * from anywhere but a stand-alone script
            if frame.f_globals["__name__"] != "__main__":
                raise RuntimeError("Cannot import * from anywhere other than \
a stand-alone script.  Do a 'from pbs import program' instead. Please see \
\"Limitations\" here: %s" % PROJECT_URL)

            warnings.warn("Importing * from pbs is magical and therefore has \
some limitations.  Please become familiar with them under \"Limitations\" \
here: %s  To avoid this warning, use a warning filter or import your \
programs directly with \"from pbs import <program>\"" % PROJECT_URL,
RuntimeWarning, stacklevel=2)

            # we avoid recursion by removing the line that imports us :)
            source = "".join(source[line:])

            exit_code = 0
            try: exec(source, env, env)
            except SystemExit as e: exit_code = e.code
            except: print(traceback.format_exc())

            # we exit so we don't actually run the script that we were imported
            # from (which would be running it "again", since we just executed
            # the script with exec
            exit(exit_code)

        # this is the least magical choice.  we're importing either a
        # selection of programs or we're just importing the pbs module.
        # in this case, let's just wrap ourselves with a module that has
        # __getattr__ so our program lookups can be done there
        else:
            self = sys.modules[__name__]
            sys.modules[__name__] = SelfWrapper(self)
<|endoftext|>


<commit_before>import math
import unittest

from numpy import array, arange, allclose

from mkufunc import Cfunc, genufunc, mkufunc

class Util:

    def assertClose(self, x, y):
        self.assert_(allclose(x, y), '%s != %s' % (x, y))

class Internal_Tests(unittest.TestCase, Util):

    def test_Cfunc(self):
        def sqr(x):
            return x * x
        cf = Cfunc(sqr, [int, int], 42)
        self.assertEqual(cf.nin, 1)
        self.assertEqual(cf.nout, 1)
        self.assertEqual(cf.cname, 'f42_pypy_g_sqr')

    def test_genufunc(self):
        def foo(x):
            return x + 17
        uf = genufunc(foo, [
                (float, float),
                (int, int),
                ])
        self.assertEqual(uf(4), 21)
        x = array([1.1, 2.3])
        y = uf(x)
        self.assertClose(y, [18.1, 19.3])
        self.assert_(str(y.dtype).startswith('float'))

        x = array([1, 4])
        y = uf(x)
        self.assertEqual(list(y), [18, 21])
        self.assert_(str(y.dtype).startswith('int'))

class Arg_Tests(unittest.TestCase, Util):

    def check_ufunc(self, f):
        for arg in (array([0.0, 1.0, 2.5]),
                    [0.0, 1.0, 2.5],
                    (0.0, 1.0, 2.5)):
            self.assertClose(f(arg), [0.0, 1.0, 6.25])

        self.assertEqual(f(3), 9)
        self.assert_(f(-2.5) - 6.25 < 1E-10)

    def test_direct(self):
        @mkufunc
        def f(x):
            return x * x
        self.check_ufunc(f)

    def test_noargs(self):
        @mkufunc()
        def f(x):
            return x * x
        self.check_ufunc(f)

    def test_varargs(self):
        for arg in (float,
                    [float],
                    [(float, float)]):
            @mkufunc(arg)
            def f(x):
                return x * x
            self.check_ufunc(f)

    def test_int(self):
        @mkufunc(int)
        def f(x):
            return x * x
        self.assertEqual(f(3), 9)
        self.assert_(isinstance(f(42), int))

    def test_mixed(self):
        @mkufunc([(int, float, int), float])
        def f(n, x):
            return n + x * x

        y = f(2, 3.9)            # Note that int(2 + 3.9 * 3.9) = 17
        self.assertEqual(y, 17)
        self.assert_(isinstance(y, int))

        y = f(2.0, 3.9)
        self.assertClose(y, 17.21)
        self.assert_(isinstance(y, float))

    def test_exceptions(self):
        def f(x):
            return x

        self.assertRaises(TypeError, mkufunc, {})
        self.assertRaises(TypeError, mkufunc([(float,)]), f)
        self.assertRaises(TypeError, mkufunc([3*(float,)]), f)
        self.assertRaises(TypeError, mkufunc([{}]), f)
        self.assertRaises(TypeError, mkufunc([(int, {})]), f)
        self.assertRaises(ValueError, mkufunc([]), f)

class Math_Tests(unittest.TestCase, Util):

    def assertFuncsEqual(self, uf, f):
        x = 0.4376
        a = uf(x)
        b = f(x)
        self.assertClose(a, b)
        xx = arange(0.1, 0.9, 0.01)
        a = uf(xx)
        b = [f(x) for x in xx]
        self.assertClose(a, b)

    def test_exp(self):
        @mkufunc
        def f(x): return math.exp(x)
        self.assertFuncsEqual(f, math.exp)

    def test_log(self):
        @mkufunc
        def f(x): return math.log(x)
        self.assertFuncsEqual(f, math.log)

    def test_sqrt(self):
        @mkufunc
        def f(x): return math.sqrt(x)
        self.assertFuncsEqual(f, math.sqrt)

    def test_cos(self):
        @mkufunc
        def f(x): return math.cos(x)
        self.assertFuncsEqual(f, math.cos)

    def test_sin(self):
        @mkufunc
        def f(x): return math.sin(x)
        self.assertFuncsEqual(f, math.sin)

    def test_tan(self):
        @mkufunc
        def f(x): return math.tan(x)
        self.assertFuncsEqual(f, math.tan)

    def test_acos(self):
        @mkufunc
        def f(x): return math.acos(x)
        self.assertFuncsEqual(f, math.acos)

    def test_asin(self):
        @mkufunc
        def f(x): return math.asin(x)
        self.assertFuncsEqual(f, math.asin)

    def test_atan(self):
        @mkufunc
        def f(x): return math.atan(x)
        self.assertFuncsEqual(f, math.atan)

    def test_atan2(self):
        @mkufunc
        def f(x, y):
            return math.atan2(x, y)

        self.assertClose(f(4, 5), math.atan2(4, 5))

        xx = array([1.0, 3.0, -2.4,  3.1, -2.3])
        yy = array([1.0, 2.0,  7.5, -8.7,  0.0])
        a = f(xx, yy)
        b = [math.atan2(x, y) for x, y in zip(xx, yy)]
        self.assertClose(a, b)

    def test_arithmetic(self):
        def f(x):
            return (4 * x + 2) / (x * x - 7 * x + 1)
        uf = mkufunc(f)
        x = arange(0, 2, 0.1)
        self.assertClose(uf(x), f(x))

        def f(x, y, z):
            return x * y * z
        uf = mkufunc(f)
        x = arange(0, 1, 0.1)
        y = 2 * x
        z = 3 * x
        self.assertClose(uf(x, y, z), f(x, y, z))

class Control_Flow_Tests(unittest.TestCase):

    def test_if(self):
        @mkufunc(int)
        def f(n):
            if n < 4:
                return n
            else:
                return n * n

        self.assertEqual(f(3), 3)
        self.assertEqual(f(4), 16)

    def test_switch(self):
        @mkufunc(int)
        def f(n):
            if n < 4:
                return n
            elif n == 4:
                return 42
            elif n == 5:
                return 73
            else:
                return n * n

        self.assertEqual(f(3), 3)
        self.assertEqual(f(4), 42)
        self.assertEqual(f(5), 73)
        self.assertEqual(f(6), 36)

    def test_loop(self):
        @mkufunc(int)
        def f(n):
            res = 0
            for i in xrange(n):
                res += i*i
            return res

        self.assertEqual(f(3), 5)
        self.assertEqual(f(95), 281295)

class FreeVariable_Tests(unittest.TestCase, Util):

    def test_const(self):
        a = 13.6
        @mkufunc
        def f(x):
            return a * x

        x = arange(0, 1, 0.1)
        self.assertClose(f(x), a * x)

    def test_const2(self):
        from math import sin, pi, sqrt
        @mkufunc
        def sin_deg(angle):
            return sin(angle / 180.0 * pi)

        self.assertClose(sin_deg([0, 30, 45, 60, 90, 180, 270, 360]),
                         [0, 0.5, 1/sqrt(2), sqrt(3)/2, 1, 0, -1, 0])

class Misc_Tests(unittest.TestCase, Util):

    pass

if __name__ == '__main__':
    unittest.main()
<commit_msg>Added test for function with no args<commit_after>import math
import unittest

from numpy import array, arange, allclose

from mkufunc import Cfunc, genufunc, mkufunc

class Util:

    def assertClose(self, x, y):
        self.assert_(allclose(x, y), '%s != %s' % (x, y))

class Internal_Tests(unittest.TestCase, Util):

    def test_Cfunc(self):
        def sqr(x):
            return x * x
        cf = Cfunc(sqr, [int, int], 42)
        self.assertEqual(cf.nin, 1)
        self.assertEqual(cf.nout, 1)
        self.assertEqual(cf.cname, 'f42_pypy_g_sqr')

    def test_genufunc(self):
        def foo(x):
            return x + 17
        uf = genufunc(foo, [
                (float, float),
                (int, int),
                ])
        self.assertEqual(uf(4), 21)
        x = array([1.1, 2.3])
        y = uf(x)
        self.assertClose(y, [18.1, 19.3])
        self.assert_(str(y.dtype).startswith('float'))

        x = array([1, 4])
        y = uf(x)
        self.assertEqual(list(y), [18, 21])
        self.assert_(str(y.dtype).startswith('int'))

class Arg_Tests(unittest.TestCase, Util):

    def check_ufunc(self, f):
        for arg in (array([0.0, 1.0, 2.5]),
                    [0.0, 1.0, 2.5],
                    (0.0, 1.0, 2.5)):
            self.assertClose(f(arg), [0.0, 1.0, 6.25])

        self.assertEqual(f(3), 9)
        self.assert_(f(-2.5) - 6.25 < 1E-10)

    def test_direct(self):
        @mkufunc
        def f(x):
            return x * x
        self.check_ufunc(f)

    def test_noargs(self):
        @mkufunc()
        def f(x):
            return x * x
        self.check_ufunc(f)

    def test_varargs(self):
        for arg in (float,
                    [float],
                    [(float, float)]):
            @mkufunc(arg)
            def f(x):
                return x * x
            self.check_ufunc(f)

    def test_int(self):
        @mkufunc(int)
        def f(x):
            return x * x
        self.assertEqual(f(3), 9)
        self.assert_(isinstance(f(42), int))

    def test_mixed(self):
        @mkufunc([(int, float, int), float])
        def f(n, x):
            return n + x * x

        y = f(2, 3.9)            # Note that int(2 + 3.9 * 3.9) = 17
        self.assertEqual(y, 17)
        self.assert_(isinstance(y, int))

        y = f(2.0, 3.9)
        self.assertClose(y, 17.21)
        self.assert_(isinstance(y, float))

    def test_exceptions(self):
        def f(x):
            return x

        self.assertRaises(TypeError, mkufunc, {})
        self.assertRaises(TypeError, mkufunc([(float,)]), f)
        self.assertRaises(TypeError, mkufunc([3*(float,)]), f)
        self.assertRaises(TypeError, mkufunc([{}]), f)
        self.assertRaises(TypeError, mkufunc([(int, {})]), f)
        self.assertRaises(ValueError, mkufunc([]), f)

class Math_Tests(unittest.TestCase, Util):

    def assertFuncsEqual(self, uf, f):
        x = 0.4376
        a = uf(x)
        b = f(x)
        self.assertClose(a, b)
        xx = arange(0.1, 0.9, 0.01)
        a = uf(xx)
        b = [f(x) for x in xx]
        self.assertClose(a, b)

    def test_exp(self):
        @mkufunc
        def f(x): return math.exp(x)
        self.assertFuncsEqual(f, math.exp)

    def test_log(self):
        @mkufunc
        def f(x): return math.log(x)
        self.assertFuncsEqual(f, math.log)

    def test_sqrt(self):
        @mkufunc
        def f(x): return math.sqrt(x)
        self.assertFuncsEqual(f, math.sqrt)

    def test_cos(self):
        @mkufunc
        def f(x): return math.cos(x)
        self.assertFuncsEqual(f, math.cos)

    def test_sin(self):
        @mkufunc
        def f(x): return math.sin(x)
        self.assertFuncsEqual(f, math.sin)

    def test_tan(self):
        @mkufunc
        def f(x): return math.tan(x)
        self.assertFuncsEqual(f, math.tan)

    def test_acos(self):
        @mkufunc
        def f(x): return math.acos(x)
        self.assertFuncsEqual(f, math.acos)

    def test_asin(self):
        @mkufunc
        def f(x): return math.asin(x)
        self.assertFuncsEqual(f, math.asin)

    def test_atan(self):
        @mkufunc
        def f(x): return math.atan(x)
        self.assertFuncsEqual(f, math.atan)

    def test_atan2(self):
        @mkufunc
        def f(x, y):
            return math.atan2(x, y)

        self.assertClose(f(4, 5), math.atan2(4, 5))

        xx = array([1.0, 3.0, -2.4,  3.1, -2.3])
        yy = array([1.0, 2.0,  7.5, -8.7,  0.0])
        a = f(xx, yy)
        b = [math.atan2(x, y) for x, y in zip(xx, yy)]
        self.assertClose(a, b)

    def test_arithmetic(self):
        def f(x):
            return (4 * x + 2) / (x * x - 7 * x + 1)
        uf = mkufunc(f)
        x = arange(0, 2, 0.1)
        self.assertClose(uf(x), f(x))

class FuncArg_Tests(unittest.TestCase, Util):

    def test_fargs0(self):
        def f():
            return 42

        uf = mkufunc(f)
        self.assertEqual(uf(), 42)
        self.assert_(type(uf()).__name__.startswith('float'))

        uf = mkufunc(int)(f)
        self.assertEqual(uf(), 42)
        self.assert_(type(uf()).__name__.startswith('int'))

class Control_Flow_Tests(unittest.TestCase):

    def test_if(self):
        @mkufunc(int)
        def f(n):
            if n < 4:
                return n
            else:
                return n * n

        self.assertEqual(f(3), 3)
        self.assertEqual(f(4), 16)

    def test_switch(self):
        @mkufunc(int)
        def f(n):
            if n < 4:
                return n
            elif n == 4:
                return 42
            elif n == 5:
                return 73
            else:
                return n * n

        self.assertEqual(f(3), 3)
        self.assertEqual(f(4), 42)
        self.assertEqual(f(5), 73)
        self.assertEqual(f(6), 36)

    def test_loop(self):
        @mkufunc(int)
        def f(n):
            res = 0
            for i in xrange(n):
                res += i*i
            return res

        self.assertEqual(f(3), 5)
        self.assertEqual(f(95), 281295)

class FreeVariable_Tests(unittest.TestCase, Util):

    def test_const(self):
        a = 13.6
        @mkufunc
        def f(x):
            return a * x

        x = arange(0, 1, 0.1)
        self.assertClose(f(x), a * x)

    def test_const2(self):
        from math import sin, pi, sqrt
        @mkufunc
        def sin_deg(angle):
            return sin(angle / 180.0 * pi)

        self.assertClose(sin_deg([0, 30, 45, 60, 90, 180, 270, 360]),
                         [0, 0.5, 1/sqrt(2), sqrt(3)/2, 1, 0, -1, 0])

class Misc_Tests(unittest.TestCase, Util):

    pass

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