<commit_before>import sublime, sublime_plugin, os, re, time

class SmartRegionOpen(sublime_plugin.TextCommand):
  def run(self, edit, **args):
    if SmartRegion.get_setting("debug"):
      sublime.active_window().run_command('show_panel', {"panel": "console", "toggle": False})
      print('>>>>>>>>>>>>>>>>>>> SmartRegionOpen Debug:')

    target = SmartRegion.get_target(self.view, args)

    if not target:
      return False

    line = False

    if re.search(':', target):
      target, line = target.split(':')

    if SmartRegion.get_setting("debug"):
      print('target:', target)
      print('line:', line)

    if target:
      if(os.path.isfile(target)):
        self.open_file(target, line)
      else:
        founded_files = []
        first_root = False
        if self.view.window() and self.view.window().project_data():
          for folder in self.view.window().project_data()['folders']:
            for root, dirs, files in os.walk(folder["path"]):
              # [TODO] Needs a better strategy...
              if re.search(r"\.git([^i]|$)", root):
                next
              else:
                if not first_root:
                  first_root = root
                for file_name in files:
                  file_path = root + '/' + file_name
                  if re.search(target, file_name) or re.search(target, file_path):
                    if(os.path.isfile(file_name)):
                      founded_files.append(file_name)
                    elif(os.path.isfile(file_path)):
                      founded_files.append(file_path)

        if SmartRegion.get_setting("debug"):
          print('founded_files:')
          print(founded_files)

        if len(founded_files) > 1:
          relative_path = target.replace(sublime.active_window().extract_variables()['folder'] + '/', '')
          self.view.window().run_command("show_overlay", {"overlay": "goto", "show_files": True, "text": relative_path})
          sublime.status_message('SmartRegion | Search for file.')
        elif len(founded_files) == 1:
          self.open_file(founded_files[0], line)
          sublime.status_message('SmartRegion | File opened.')
        else:
          sublime.status_message('SmartRegion | File not found: ' + target + ' | ' + first_root)

  def go_to_line(self, view, line):
    if not view.is_loading():
      line = int(line)

      if SmartRegion.get_setting("debug"):
        print('go to line: view loaded!')
        print('go to line:', line)

      view.sel().clear()
      view.sel().add(sublime.Region(view.text_point(line-1, 0)))
      view.show(view.text_point(line-1, 0))
    else:
      if SmartRegion.get_setting("debug"):
        print('go to line: view not loaded yet...')
      sublime.set_timeout(lambda: self.go_to_line(view, line), 10)

  def open_file(self, file, line):
    if SmartRegion.get_setting("debug"):
      print('open_file:', file)

    file_view = self.view.window().find_open_file(file)

    if file_view:
      if SmartRegion.get_setting("debug"):
        print('view exists:', 'focus')
      self.view.window().focus_view(file_view)
    else:
      if SmartRegion.get_setting("debug"):
        print('view not exists:', 'open')
      file_view = self.view.window().open_file(file)

    if line:
      self.go_to_line(file_view, line)

  def want_event(self):
    return True

class SmartRegionCreateRegions(sublime_plugin.TextCommand):
  def run(self, edit, **args):
    SmartRegion.create_regions(self.view)

  def want_event(self):
    return True

class SmartRegion():
  def get_target(view, args):
    target = False
    sel_target = False

    if args.get("event"):
      click_text_position = view.window_to_text((args["event"]["x"], args["event"]["y"]))
    else:
      sel_target = view.substr(view.sel()[0])
      click_text_position = view.sel()[0].begin()

    for region in view.get_regions('smart_region'):
      if region.contains(click_text_position):
        target = view.substr(region)

    if not target:
      target = sel_target

    if not target:
      sublime.status_message('SmartRegion | No Target')
    else:
      return target

  def create_regions(view):
    if view.size() > 100000:
      sublime.status_message('SmartRegion | File is too large: ' + str(view.size()) + ' chars')
    else:
      regions = []

      if view.window() and view.window().project_data():
        for folder in view.window().project_data()['folders']:
          for root, dirs, files in os.walk(folder["path"]):

            # [TODO] Needs a better strategy...
            if re.search(r"\.git([^i]|$)", root):
              next
            else:
              for target in view.find_all(re.escape(root)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              for target in view.find_all(re.escape(root) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              # [TODO] Maybe this is not a good idea...
              for dir_name in dirs:
                for target in view.find_all(re.escape(dir_name)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(dir_name) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              for file_name in files:
                full_path = root + '/' + file_name
                relative_path = full_path.replace(sublime.active_window().extract_variables()['folder'] + '/', '')

                for target in view.find_all(re.escape(file_name)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(file_name) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(relative_path)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(relative_path) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(full_path)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(full_path) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

      file_regex = "((\s\/|^\/)[\w|\-|\/|\@|\.]{1,}[\.[\w|\-]{1,}|$])"

      for target in view.find_all(file_regex):
        regions.append(view.find(view.substr(target).strip(), target.begin()))

      for target in view.find_all(file_regex + ':\d{1,}'):
        regions.append(view.find(view.substr(target).strip(), target.begin()))

      flags = 0

      if "draw_empty" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_EMPTY

      if "hide_on_minimap" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.HIDE_ON_MINIMAP

      if "draw_empty_as_overwrite" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_EMPTY_AS_OVERWRITE

      if "draw_no_fill" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_NO_FILL

      if "draw_no_outline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_NO_OUTLINE

      if "draw_solid_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_SOLID_UNDERLINE

      if "draw_stippled_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_STIPPLED_UNDERLINE

      if "draw_squiggly_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_SQUIGGLY_UNDERLINE

      if "persistent" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.PERSISTENT

      if "hidden" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.HIDDEN

      view.add_regions('smart_region', regions, 'code', '', flags)

  def get_setting(config):
    settings = sublime.load_settings('Preferences.sublime-settings')
    if settings.get('smart_region_' + config):
      return settings.get('smart_region_' + config)
    else:
      settings = sublime.load_settings('SmartRegion.sublime-settings')
      return settings.get('smart_region_' + config)

class KeyBindingListener(sublime_plugin.EventListener):

  def on_load_async(self, view):
    if "load" in SmartRegion.get_setting("create_regions_on"):
      SmartRegion.create_regions(view)

  def on_modified_async(self, view):
    if "modified" in SmartRegion.get_setting("create_regions_on"):
      SmartRegion.create_regions(view)

<commit_msg>fix open search files<commit_after>import sublime, sublime_plugin, os, re, time

class SmartRegionOpen(sublime_plugin.TextCommand):
  def run(self, edit, **args):
    if SmartRegion.get_setting("debug"):
      sublime.active_window().run_command('show_panel', {"panel": "console", "toggle": False})
      print('>>>>>>>>>>>>>>>>>>> SmartRegionOpen Debug:')

    target = SmartRegion.get_target(self.view, args)

    if not target:
      return False

    line = False

    if re.search(':', target):
      target, line = target.split(':')

    if SmartRegion.get_setting("debug"):
      print('target:', target)
      print('line:', line)

    if target:
      if(os.path.isfile(target)):
        self.open_file(target, line)
      else:
        founded_files = []
        first_root = False
        if self.view.window() and self.view.window().project_data():
          for folder in self.view.window().project_data()['folders']:
            for root, dirs, files in os.walk(folder["path"]):
              # [TODO] Needs a better strategy...
              if re.search(r"\.git([^i]|$)", root):
                next
              else:
                if not first_root:
                  first_root = root
                for file_name in files:
                  file_path = root + '/' + file_name
                  if re.search(target, file_name) or re.search(target, file_path):
                    if(os.path.isfile(file_name)):
                      founded_files.append(file_name)
                    elif(os.path.isfile(file_path)):
                      founded_files.append(file_path)

        if SmartRegion.get_setting("debug"):
          print('founded_files:')
          print(founded_files)

        reversed(founded_files)

        if len(founded_files) > 1:
          relative_path = target.replace(sublime.active_window().extract_variables()['folder'] + '/', '')
          self.view.window().run_command("show_overlay", {"overlay": "goto", "show_files": True, "text": relative_path})
          sublime.status_message('SmartRegion | Search for file.')
        elif len(founded_files) == 1:
          self.open_file(founded_files[0], line)
          sublime.status_message('SmartRegion | File opened.')
        else:
          sublime.status_message('SmartRegion | File not found: ' + target + ' | ' + first_root)

  def go_to_line(self, view, line):
    if not view.is_loading():
      line = int(line)

      if SmartRegion.get_setting("debug"):
        print('go to line: view loaded!')
        print('go to line:', line)

      view.sel().clear()
      view.sel().add(sublime.Region(view.text_point(line-1, 0)))
      view.show(view.text_point(line-1, 0))
    else:
      if SmartRegion.get_setting("debug"):
        print('go to line: view not loaded yet...')
      sublime.set_timeout(lambda: self.go_to_line(view, line), 10)

  def open_file(self, file, line):
    if SmartRegion.get_setting("debug"):
      print('open_file:', file)

    file_view = self.view.window().find_open_file(file)

    if file_view:
      if SmartRegion.get_setting("debug"):
        print('view exists:', 'focus')
      self.view.window().focus_view(file_view)
    else:
      if SmartRegion.get_setting("debug"):
        print('view not exists:', 'open')
      file_view = self.view.window().open_file(file)

    if line:
      self.go_to_line(file_view, line)

  def want_event(self):
    return True

class SmartRegionCreateRegions(sublime_plugin.TextCommand):
  def run(self, edit, **args):
    SmartRegion.create_regions(self.view)

  def want_event(self):
    return True

class SmartRegion():
  def get_target(view, args):
    target = False
    sel_target = False

    if args.get("event"):
      click_text_position = view.window_to_text((args["event"]["x"], args["event"]["y"]))
    else:
      sel_target = view.substr(view.sel()[0])
      click_text_position = view.sel()[0].begin()

    for region in view.get_regions('smart_region'):
      if region.contains(click_text_position):
        target = view.substr(region)

    if not target:
      target = sel_target

    if not target:
      sublime.status_message('SmartRegion | No Target')
    else:
      return target

  def create_regions(view):
    if view.size() > 100000:
      sublime.status_message('SmartRegion | File is too large: ' + str(view.size()) + ' chars')
    else:
      regions = []

      if view.window() and view.window().project_data():
        for folder in view.window().project_data()['folders']:
          for root, dirs, files in os.walk(folder["path"]):

            # [TODO] Needs a better strategy...
            if re.search(r"\.git([^i]|$)", root):
              next
            else:
              for target in view.find_all(re.escape(root)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              for target in view.find_all(re.escape(root) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              # [TODO] Maybe this is not a good idea...
              for dir_name in dirs:
                for target in view.find_all(re.escape(dir_name)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(dir_name) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

              for file_name in files:
                full_path = root + '/' + file_name
                relative_path = full_path.replace(sublime.active_window().extract_variables()['folder'] + '/', '')

                for target in view.find_all(re.escape(file_name)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(file_name) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(relative_path)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(relative_path) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(full_path)):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

                for target in view.find_all(re.escape(full_path) + ':\d{1,}'):
                  regions.append(view.find(re.escape(view.substr(target).strip()), target.begin()))

      file_regex = "((\s\/|^\/)[\w|\-|\/|\@|\.]{1,}[\.[\w|\-]{1,}|$])"

      for target in view.find_all(file_regex):
        regions.append(view.find(view.substr(target).strip(), target.begin()))

      for target in view.find_all(file_regex + ':\d{1,}'):
        regions.append(view.find(view.substr(target).strip(), target.begin()))

      flags = 0

      if "draw_empty" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_EMPTY

      if "hide_on_minimap" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.HIDE_ON_MINIMAP

      if "draw_empty_as_overwrite" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_EMPTY_AS_OVERWRITE

      if "draw_no_fill" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_NO_FILL

      if "draw_no_outline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_NO_OUTLINE

      if "draw_solid_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_SOLID_UNDERLINE

      if "draw_stippled_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_STIPPLED_UNDERLINE

      if "draw_squiggly_underline" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.DRAW_SQUIGGLY_UNDERLINE

      if "persistent" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.PERSISTENT

      if "hidden" in SmartRegion.get_setting("draw_regions"):
        flags += sublime.HIDDEN

      view.add_regions('smart_region', regions, 'code', '', flags)

  def get_setting(config):
    settings = sublime.load_settings('Preferences.sublime-settings')
    if settings.get('smart_region_' + config):
      return settings.get('smart_region_' + config)
    else:
      settings = sublime.load_settings('SmartRegion.sublime-settings')
      return settings.get('smart_region_' + config)

class KeyBindingListener(sublime_plugin.EventListener):

  def on_load_async(self, view):
    if "load" in SmartRegion.get_setting("create_regions_on"):
      SmartRegion.create_regions(view)

  def on_modified_async(self, view):
    if "modified" in SmartRegion.get_setting("create_regions_on"):
      SmartRegion.create_regions(view)

<|endoftext|>


<commit_before>import json
import logging

from django.conf import settings
from django.contrib.auth import authenticate, login, logout
from django.contrib.auth.views import password_reset
from django.contrib.sites.models import Site
from django.contrib.sites.shortcuts import get_current_site
from django.core.exceptions import ObjectDoesNotExist
from django.core.mail import send_mail
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponse, HttpResponseBadRequest
from django.shortcuts import render, redirect
from django.template.loader import render_to_string
from django.utils.translation import ugettext as _
from django.views.decorators.csrf import csrf_exempt
from django.views.generic.edit import FormView
from registration.backends.hmac.views import RegistrationView

from contentcuration.api import add_editor_to_channel
from contentcuration.forms import RegistrationForm, RegistrationInformationForm, USAGES
from contentcuration.models import Channel, User, Invitation
from contentcuration.statistics import record_user_registration_stats

""" REGISTRATION/INVITATION ENDPOINTS """

@csrf_exempt
def send_invitation_email(request):
    if request.method != 'POST':
        raise HttpResponseBadRequest("Only POST requests are allowed on this endpoint.")
    else:
        data = json.loads(request.body)

        try:
            user_email = data["user_email"]
            channel_id = data["channel_id"]
            share_mode = data["share_mode"]
            retrieved_user = User.objects.get_or_create(email=user_email)
            recipient = retrieved_user[0]
            channel = Channel.objects.get(id=channel_id)

            request.user.can_view(channel_id)

            fields = {
                "invited": recipient,
                "email": user_email,
                "channel_id": channel_id,
                "first_name": recipient.first_name if recipient.is_active else "Guest",
                "last_name": recipient.last_name if recipient.is_active else " "
            }

            # Need to break into two steps to avoid MultipleObjectsReturned error
            invitation = Invitation.objects.filter(**fields).first()

            if not invitation:
                invitation = Invitation.objects.create(**fields)

            # Handle these values separately as different users might invite the same user again
            invitation.share_mode = share_mode
            invitation.sender = invitation.sender or request.user
            invitation.save()

            ctx_dict = {'sender': request.user,
                        'site': get_current_site(request),
                        'user': recipient,
                        'share_mode': _(share_mode),
                        'channel_id': channel_id,
                        'invitation_key': invitation.id,
                        'is_new': recipient.is_active is False,
                        'channel': channel.name,
                        'domain': request.META.get('HTTP_ORIGIN') or "http://{}".format(
                            request.get_host() or Site.objects.get_current().domain),
                        }
            subject = render_to_string('permissions/permissions_email_subject.txt', ctx_dict)
            message = render_to_string('permissions/permissions_email.txt', ctx_dict)
            # message_html = render_to_string('permissions/permissions_email.html', ctx_dict)
            recipient.email_user(subject, message, settings.DEFAULT_FROM_EMAIL, )  # html_message=message_html,)
            # recipient.email_user(subject, message, settings.DEFAULT_FROM_EMAIL,)
        except KeyError:
            raise ObjectDoesNotExist("Missing attribute from data: {}".format(data))

        return HttpResponse(json.dumps({
            "id": invitation.pk,
            "invited": invitation.invited_id,
            "email": invitation.email,
            "sender": invitation.sender_id,
            "channel": invitation.channel_id,
            "first_name": invitation.first_name,
            "last_name": invitation.last_name,
            "share_mode": invitation.share_mode,
        }))

class UserRegistrationView(RegistrationView):
    form_class = RegistrationForm

    def get_initial(self):
        initial = self.initial.copy()
        return {
            'email': self.request.session.get('email', None),
            'first_name': self.request.session.get('first_name', None),
            'last_name': self.request.session.get('last_name', None),
            'password1': self.request.session.get('password1', None),
            'password2': self.request.session.get('password2', None),
        }

    def get_context_data(self, **kwargs):
        kwargs = super(UserRegistrationView, self).get_context_data(**kwargs)
        kwargs.update({"freeze_email": self.request.session.get("freeze_email"),})
        return kwargs

    def post(self, request):
        form = self.form_class(request.POST)
        if form.is_valid():
            # Store information in session in case user goes back
            self.request.session.update(form.cleaned_data)
            return redirect(reverse_lazy('registration_information'))
        return super(UserRegistrationView, self).post(request)

class InformationRegistrationView(RegistrationView):
    email_body_template = 'registration/activation_email.txt'
    email_subject_template = 'registration/activation_email_subject.txt'
    email_html_template = 'registration/activation_email.html'
    template_name = 'registration/registration_information_form.html'
    form_class = RegistrationInformationForm

    def get_form_kwargs(self):
        kw = super(InformationRegistrationView, self).get_form_kwargs()
        kw['request'] = self.request
        return kw

    def get_context_data(self, **kwargs):
        kwargs = super(InformationRegistrationView, self).get_context_data(**kwargs)
        kwargs.update({"help_email": settings.HELP_EMAIL})
        return kwargs

    def get_initial(self):
        initial = self.initial.copy()
        return {
            'email': self.request.session.get('email', None),
            'first_name': self.request.session.get('first_name', None),
            'last_name': self.request.session.get('last_name', None),
            'password1': self.request.session.get('password1', None),
            'password2': self.request.session.get('password2', None),
        }

    def register(self, form):
        # Send email regarding new user information
        subject = render_to_string('registration/custom_email_subject.txt', {"subject": "New Kolibri Studio Registration"})
        message = render_to_string('registration/registration_information_email.txt', form.cleaned_data)
        send_mail(subject, message, settings.DEFAULT_FROM_EMAIL, [settings.REGISTRATION_INFORMATION_EMAIL])

        # Clear session cached fields
        self.request.session["freeze_email"] = False
        for field in RegistrationForm.Meta.fields:
            self.request.session[field] = ""

        return super(InformationRegistrationView, self).register(form)

    def send_activation_email(self, user):
        activation_key = self.get_activation_key(user)
        context = self.get_email_context(activation_key)
        context.update({
            'user': user,
            'domain': self.request.META.get('HTTP_ORIGIN') or "http://{}".format(
                self.request.get_host() or Site.objects.get_current().domain),
        })
        subject = render_to_string(self.email_subject_template, context)
        subject = ''.join(subject.splitlines())
        message = render_to_string(self.email_body_template, context)
        # message_html = render_to_string(self.email_html_template, context)
        user.email_user(subject, message, settings.DEFAULT_FROM_EMAIL, )  # html_message=message_html,)

        record_user_registration_stats(user)

def custom_password_reset(request, **kwargs):
    email_context = {'site': get_current_site(request), 'domain': request.META.get('HTTP_ORIGIN') or "http://{}".format(
        request.get_host() or Site.objects.get_current().domain)}
    return password_reset(request, extra_email_context=email_context, **kwargs)

def new_user_redirect(request, user_id):
    user = User.objects.get(pk=user_id)
    if user.is_active:
        return redirect(reverse_lazy("channels"))
    request.session["email"] = user.email
    request.session["freeze_email"] = True
    logout(request)
    return redirect(reverse_lazy("registration_register"))
<commit_msg>Dummy commit to trigger build<commit_after>import json
import logging

from django.conf import settings
from django.contrib.auth import authenticate, login, logout
from django.contrib.auth.views import password_reset
from django.contrib.sites.models import Site
from django.contrib.sites.shortcuts import get_current_site
from django.core.exceptions import ObjectDoesNotExist
from django.core.mail import send_mail
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponse, HttpResponseBadRequest
from django.shortcuts import render, redirect
from django.template.loader import render_to_string
from django.utils.translation import ugettext as _
from django.views.decorators.csrf import csrf_exempt
from django.views.generic.edit import FormView
from registration.backends.hmac.views import RegistrationView

from contentcuration.api import add_editor_to_channel
from contentcuration.forms import RegistrationForm, RegistrationInformationForm, USAGES
from contentcuration.models import Channel, User, Invitation
from contentcuration.statistics import record_user_registration_stats

""" REGISTRATION/INVITATION ENDPOINTS """

@csrf_exempt
def send_invitation_email(request):
    if request.method != 'POST':
        raise HttpResponseBadRequest("Only POST requests are allowed on this endpoint.")
    else:
        data = json.loads(request.body)

        try:
            user_email = data["user_email"]
            channel_id = data["channel_id"]
            share_mode = data["share_mode"]
            retrieved_user = User.objects.get_or_create(email=user_email)
            recipient = retrieved_user[0]
            channel = Channel.objects.get(id=channel_id)

            request.user.can_view(channel_id)

            fields = {
                "invited": recipient,
                "email": user_email,
                "channel_id": channel_id,
                "first_name": recipient.first_name if recipient.is_active else "Guest",
                "last_name": recipient.last_name if recipient.is_active else " "
            }

            # Need to break into two steps to avoid MultipleObjectsReturned error
            invitation = Invitation.objects.filter(**fields).first()

            if not invitation:
                invitation = Invitation.objects.create(**fields)

            # Handle these values separately as different users might invite the same user again
            invitation.share_mode = share_mode
            invitation.sender = invitation.sender or request.user
            invitation.save()

            ctx_dict = {'sender': request.user,
                        'site': get_current_site(request),
                        'user': recipient,
                        'share_mode': _(share_mode),
                        'channel_id': channel_id,
                        'invitation_key': invitation.id,
                        'is_new': recipient.is_active is False,
                        'channel': channel.name,
                        'domain': request.META.get('HTTP_ORIGIN') or "http://{}".format(
                            request.get_host() or Site.objects.get_current().domain),
                        }
            subject = render_to_string('permissions/permissions_email_subject.txt', ctx_dict)
            message = render_to_string('permissions/permissions_email.txt', ctx_dict)
            # message_html = render_to_string('permissions/permissions_email.html', ctx_dict)
            recipient.email_user(subject, message, settings.DEFAULT_FROM_EMAIL, )  # html_message=message_html,)
            # recipient.email_user(subject, message, settings.DEFAULT_FROM_EMAIL,)
        except KeyError:
            raise ObjectDoesNotExist("Missing attribute from data: {}".format(data))

        return HttpResponse(json.dumps({
            "id": invitation.pk,
            "invited": invitation.invited_id,
            "email": invitation.email,
            "sender": invitation.sender_id,
            "channel": invitation.channel_id,
            "first_name": invitation.first_name,
            "last_name": invitation.last_name,
            "share_mode": invitation.share_mode,
        }))

class UserRegistrationView(RegistrationView):
    form_class = RegistrationForm

    def get_initial(self):
        initial = self.initial.copy()
        return {
            'email': self.request.session.get('email', None),
            'first_name': self.request.session.get('first_name', None),
            'last_name': self.request.session.get('last_name', None),
            'password1': self.request.session.get('password1', None),
            'password2': self.request.session.get('password2', None),
        }

    def get_context_data(self, **kwargs):
        kwargs = super(UserRegistrationView, self).get_context_data(**kwargs)
        kwargs.update({"freeze_email": self.request.session.get("freeze_email"),})
        return kwargs

    def post(self, request):
        form = self.form_class(request.POST)
        if form.is_valid():
            # Store information in session in case user goes back
            self.request.session.update(form.cleaned_data)
            return redirect(reverse_lazy('registration_information'))
        return super(UserRegistrationView, self).post(request)

class InformationRegistrationView(RegistrationView):
    email_body_template = 'registration/activation_email.txt'
    email_subject_template = 'registration/activation_email_subject.txt'
    email_html_template = 'registration/activation_email.html'
    template_name = 'registration/registration_information_form.html'
    form_class = RegistrationInformationForm

    def get_form_kwargs(self):
        kw = super(InformationRegistrationView, self).get_form_kwargs()
        kw['request'] = self.request
        return kw

    def get_context_data(self, **kwargs):
        kwargs = super(InformationRegistrationView, self).get_context_data(**kwargs)
        kwargs.update({"help_email": settings.HELP_EMAIL})
        return kwargs

    def get_initial(self):
        initial = self.initial.copy()
        return {
            'email': self.request.session.get('email', None),
            'first_name': self.request.session.get('first_name', None),
            'last_name': self.request.session.get('last_name', None),
            'password1': self.request.session.get('password1', None),
            'password2': self.request.session.get('password2', None),
        }

    def register(self, form):
        # Send email regarding new user information
        subject = render_to_string('registration/custom_email_subject.txt', {"subject": "New Kolibri Studio Registration"})
        message = render_to_string('registration/registration_information_email.txt', form.cleaned_data)
        send_mail(subject, message, settings.DEFAULT_FROM_EMAIL, [settings.REGISTRATION_INFORMATION_EMAIL])

        # Clear session cached fields
        self.request.session["freeze_email"] = False
        for field in RegistrationForm.Meta.fields:
            self.request.session[field] = ""

        return super(InformationRegistrationView, self).register(form)

    def send_activation_email(self, user):
        activation_key = self.get_activation_key(user)
        context = self.get_email_context(activation_key)
        context.update({
            'user': user,
            'domain': self.request.META.get('HTTP_ORIGIN') or "http://{}".format(
                self.request.get_host() or Site.objects.get_current().domain),
        })
        subject = render_to_string(self.email_subject_template, context)
        subject = ''.join(subject.splitlines())
        message = render_to_string(self.email_body_template, context)
        # message_html = render_to_string(self.email_html_template, context)
        user.email_user(subject, message, settings.DEFAULT_FROM_EMAIL, )  # html_message=message_html,)

        record_user_registration_stats(user)

def custom_password_reset(request, **kwargs):
    email_context = {'site': get_current_site(request), 'domain': request.META.get('HTTP_ORIGIN') or "http://{}".format(
        request.get_host() or Site.objects.get_current().domain)}
    return password_reset(request, extra_email_context=email_context, **kwargs)

def new_user_redirect(request, user_id):
    user = User.objects.get(pk=user_id)
    if user.is_active:
        return redirect(reverse_lazy("channels"))
    request.session["email"] = user.email
    request.session["freeze_email"] = True
    logout(request)
    return redirect(reverse_lazy("registration_register"))
<|endoftext|>


<commit_before><commit_msg>Revert of Disable testWebpageReplay to green the cros perf bots. (patchset #1 id:1 of https://codereview.chromium.org/1250363007/ )<commit_after><|endoftext|>


<commit_before># Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

"""Initial revision

Revision ID: f16eb4eff7c
Revises:
Create Date: 2016-03-02 17:10:04.750584

"""

# revision identifiers, used by Alembic.
revision = 'f16eb4eff7c'
down_revision = None
branch_labels = None
depends_on = None

from alembic import context
from alembic import op
import sqlalchemy as sa

import tuning_box.db

def upgrade():
    table_prefix = context.config.get_main_option('table_prefix')
    op.create_table(
        table_prefix + 'component',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
    )
    op.create_table(
        table_prefix + 'environment',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
    )
    op.create_table(
        table_prefix + 'namespace',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
    )
    op.create_table(
        table_prefix + 'environment_components',
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id']),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id']),
    )
    op.create_table(
        table_prefix + 'environment_hierarchy_level',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('parent_id', sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id']),
        sa.ForeignKeyConstraint(
            ['parent_id'], [table_prefix + 'environment_hierarchy_level.id']),
        sa.UniqueConstraint('environment_id', 'name'),
        sa.UniqueConstraint('environment_id', 'parent_id'),
    )
    op.create_table(
        table_prefix + 'schema',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.Column('namespace_id', sa.Integer(), nullable=True),
        sa.Column('content', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id']),
        sa.ForeignKeyConstraint(
            ['namespace_id'], [table_prefix + 'namespace.id']),
    )
    op.create_table(
        table_prefix + 'template',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.Column('content', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id']),
    )
    op.create_table(
        table_prefix + 'environment_hierarchy_level_value',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('level_id', sa.Integer(), nullable=True),
        sa.Column('parent_id', sa.Integer(), nullable=True),
        sa.Column('value', sa.String(length=128), nullable=True),
        sa.ForeignKeyConstraint(
            ['level_id'], [table_prefix + 'environment_hierarchy_level.id']),
        sa.ForeignKeyConstraint(
            ['parent_id'],
            [table_prefix + 'environment_hierarchy_level_value.id']),
    )
    op.create_table(
        table_prefix + 'environment_schema_values',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('schema_id', sa.Integer(), nullable=True),
        sa.Column('level_value_id', sa.Integer(), nullable=True),
        sa.Column('values', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id']),
        sa.ForeignKeyConstraint(
            ['level_value_id'],
            [table_prefix + 'environment_hierarchy_level_value.id']),
        sa.ForeignKeyConstraint(['schema_id'], [table_prefix + 'schema.id']),
        sa.UniqueConstraint('environment_id', 'schema_id', 'level_value_id'),
    )

def downgrade():
    table_prefix = context.config.get_main_option('table_prefix')
    op.drop_table(table_prefix + 'environment_schema_values')
    op.drop_table(table_prefix + 'environment_hierarchy_level_value')
    op.drop_table(table_prefix + 'template')
    op.drop_table(table_prefix + 'schema')
    op.drop_table(table_prefix + 'environment_hierarchy_level')
    op.drop_table(table_prefix + 'environment_components')
    op.drop_table(table_prefix + 'namespace')
    op.drop_table(table_prefix + 'environment')
    op.drop_table(table_prefix + 'component')
<commit_msg>Add foreign key names to all FKs in initial migration<commit_after># Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

"""Initial revision

Revision ID: f16eb4eff7c
Revises:
Create Date: 2016-03-02 17:10:04.750584

"""

# revision identifiers, used by Alembic.
revision = 'f16eb4eff7c'
down_revision = None
branch_labels = None
depends_on = None

from alembic import context
from alembic import op
import sqlalchemy as sa

import tuning_box.db

def upgrade():
    table_prefix = context.config.get_main_option('table_prefix')
    op.create_table(
        table_prefix + 'component',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
    )
    op.create_table(
        table_prefix + 'environment',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
    )
    op.create_table(
        table_prefix + 'namespace',
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
    )
    table_name = table_prefix + 'environment_components'
    op.create_table(
        table_name,
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id'],
            name=table_name + '_component_id_fkey',
        ),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id'],
            name=table_name + '_environment_id_fkey',
        ),
    )
    table_name = table_prefix + 'environment_hierarchy_level'
    op.create_table(
        table_name,
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('parent_id', sa.Integer(), nullable=True),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id'],
            name=table_name + '_environment_id_fkey',
        ),
        sa.ForeignKeyConstraint(
            ['parent_id'], [table_prefix + 'environment_hierarchy_level.id'],
            name=table_name + '_parent_id_fkey',
        ),
        sa.UniqueConstraint(
            'environment_id', 'name',
            name=table_name + '_environment_id_name_key',
        ),
        sa.UniqueConstraint(
            'environment_id', 'parent_id',
            name=table_name[:-4] + '_environment_id_parent_id_key',
        ),
    )
    table_name = table_prefix + 'schema'
    op.create_table(
        table_name,
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.Column('namespace_id', sa.Integer(), nullable=True),
        sa.Column('content', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id'],
            name=table_name + '_component_id_fkey'),
        sa.ForeignKeyConstraint(
            ['namespace_id'], [table_prefix + 'namespace.id'],
            name=table_name + '_namespace_id_fkey'),
    )
    table_name = table_prefix + 'template'
    op.create_table(
        table_name,
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('name', sa.String(length=128), nullable=True),
        sa.Column('component_id', sa.Integer(), nullable=True),
        sa.Column('content', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['component_id'], [table_prefix + 'component.id'],
            name=table_name + '_component_id_fkey',
        ),
    )
    table_name = table_prefix + 'environment_hierarchy_level_value'
    op.create_table(
        table_name,
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('level_id', sa.Integer(), nullable=True),
        sa.Column('parent_id', sa.Integer(), nullable=True),
        sa.Column('value', sa.String(length=128), nullable=True),
        sa.ForeignKeyConstraint(
            ['level_id'], [table_prefix + 'environment_hierarchy_level.id'],
            name=table_name + '_level_id_fkey',
        ),
        sa.ForeignKeyConstraint(
            ['parent_id'], [table_name + '.id'],
            name=table_name + '_parent_id_fkey',
        ),
    )
    table_name = table_prefix + 'environment_schema_values'
    op.create_table(
        table_name,
        sa.Column('id', sa.Integer(), nullable=False, primary_key=True),
        sa.Column('environment_id', sa.Integer(), nullable=True),
        sa.Column('schema_id', sa.Integer(), nullable=True),
        sa.Column('level_value_id', sa.Integer(), nullable=True),
        sa.Column('values', tuning_box.db.Json(), nullable=True),
        sa.ForeignKeyConstraint(
            ['environment_id'], [table_prefix + 'environment.id'],
            name=table_name + '_environment_id_fkey',
        ),
        sa.ForeignKeyConstraint(
            ['level_value_id'],
            [table_prefix + 'environment_hierarchy_level_value.id'],
            name=table_name + '_level_value_id_fkey',
        ),
        sa.ForeignKeyConstraint(
            ['schema_id'], [table_prefix + 'schema.id'],
            name=table_name + '_schema_id_fkey',
        ),
        sa.UniqueConstraint(
            'environment_id', 'schema_id', 'level_value_id',
            name=table_name[:-6] + 'environment_id_schema_id_leve_key',
        ),
    )

def downgrade():
    table_prefix = context.config.get_main_option('table_prefix')
    op.drop_table(table_prefix + 'environment_schema_values')
    op.drop_table(table_prefix + 'environment_hierarchy_level_value')
    op.drop_table(table_prefix + 'template')
    op.drop_table(table_prefix + 'schema')
    op.drop_table(table_prefix + 'environment_hierarchy_level')
    op.drop_table(table_prefix + 'environment_components')
    op.drop_table(table_prefix + 'namespace')
    op.drop_table(table_prefix + 'environment')
    op.drop_table(table_prefix + 'component')
<|endoftext|>


<commit_before>import wee_slack
import pytest
import json

slack = wee_slack

unfurl_map = [
    { "input": "foo",
      "output": "foo",
    },
    { "input": "<@U2147483697|@othernick>: foo",
      "output": "@testuser: foo",
      "ignore_alt_text": True
    },
    { "input": "foo <#C2147483705|#otherchannel> foo",
      "output": "foo #otherchannel foo",
    },
    { "input": "foo <#C2147483705> foo",
      "output": "foo #testchan foo",
    },
    { "input": "url: <https://example.com|example> suffix",
      "output": "url: https://example.com (example) suffix",
    },
    { "input": "url: <https://example.com|example with spaces> suffix",
      "output": "url: https://example.com (example with spaces) suffix",
    },
    { "input": "<@U2147483697|@othernick> multiple unfurl <https://example.com|example with spaces>",
      "output": "@othernick multiple unfurl https://example.com (example with spaces)",
    },
    ]

def test_unfurl_refs(myservers, mychannels, myusers):
    slack.servers = myservers
    slack.channels = mychannels
    slack.users = myusers
    slack.message_cache = {}
    slack.servers[0].users = myusers
    print mychannels[0].identifier

    for k in unfurl_map:
        if "ignore_alt_text" in k:
            assert slack.unfurl_refs(k["input"], ignore_alt_text=k["ignore_alt_text"]) == k["output"]
        else:
            assert slack.unfurl_refs(k["input"]) == k["output"]
<commit_msg>Leverage pytest.parametrize for test_unfurl<commit_after>import wee_slack
import pytest

slack = wee_slack

@pytest.mark.parametrize('case', (
    {
        'input': "foo",
        'output': "foo",
    },
    {
        'input': "<@U2147483697|@othernick>: foo",
        'output': "@testuser: foo",
        'ignore_alt_text': True,
    },
    {
        'input': "foo <#C2147483705|#otherchannel> foo",
        'output': "foo #otherchannel foo",
    },
    {
        'input': "foo <#C2147483705> foo",
        'output': "foo #testchan foo",
    },
    {
        'input': "url: <https://example.com|example> suffix",
        'output': "url: https://example.com (example) suffix",
    },
    {
        'input': "url: <https://example.com|example with spaces> suffix",
        'output': "url: https://example.com (example with spaces) suffix",
    },
    {
        'input': "<@U2147483697|@othernick> multiple unfurl <https://example.com|example with spaces>",
        'output': "@othernick multiple unfurl https://example.com (example with spaces)",
    },
))
def test_unfurl_refs(myservers, mychannels, myusers, case):
    slack.servers = myservers
    slack.channels = mychannels
    slack.users = myusers
    slack.message_cache = {}
    slack.servers[0].users = myusers
    print mychannels[0].identifier

    assert slack.unfurl_refs(case['input'], ignore_alt_text=case.get('ignore_alt_text', False)) == case['output']
<|endoftext|>


<commit_before>
from operator import itemgetter

from django.db.models import Count
from django.db.models import Sum
from django.db.models import Q, F
from django.db.models.functions import Coalesce
from django.db import connection

def apply_annotations(queryset, selected_groupings, selected_aggregations, query_params):
    """
    Builds and performs the query, when multiple aggregations were requested it joins the results
    """
    result_dict = None

    #
    # Apply group_by fields and renames
    #

    group_fields = flatten([ grouping.get_fields() for grouping in selected_groupings ])
    rename_annotations = merge([ grouping.get_renamed_fields() for grouping in selected_groupings ])
    group_extras = merge([ grouping.extra for grouping in selected_groupings if grouping.extra is not None])

    eliminate_nulls = {"{}__isnull".format(grouping): False for grouping in group_fields}

    queryset = queryset \
        .extra(**group_extras) \
        .annotate(**rename_annotations)

    if len(group_extras) is 0:
        queryset = queryset \
            .filter(**eliminate_nulls)

    # preparation for aggregation look
    main_group_key = group_fields[0]
    rest_group_keys = group_fields[1:]

    def get_aggregation_queryset(queryset, group_fields, aggregation):

        # TODO: Should queryset be copied here? - 2016-04-07
        next_result = queryset.all()

        # apply any extra aggregation filters if specified
        next_result = aggregation.apply_extra_filter(next_result)

        # apply group_by values() call
        next_result = next_result.values(*group_fields)

        # apply the aggregation annotation
        next_result = aggregation.apply_annotation(next_result, query_params, selected_groupings)
        # print str(next_result.query)
        return next_result

    aggregation_querysets = [
        get_aggregation_queryset(queryset, group_fields, aggregation)
            for aggregation in selected_aggregations
    ]

    def merge_results(querysets, group_fields):
        """
        Execute the querysets and merge the results into one list of dictionaries
        This method keeps ordering of keys in order of execution of the aggregations
        """

        if len(querysets) is 1:
            return querysets[0]

        first_queryset = querysets[0]
        next_querysets = querysets[1:]

        result_dict = {}

        def set_item_keys(item, group_fields):

            for aggregation in selected_aggregations:
                if not aggregation.field in item:
                    item[aggregation.field] = 0
            return item

        def get_group_key(item, group_fields):
            """
            Generate a unique group key from fields that are grouped by
            """
            if len(group_fields) is 1:
                return item[group_fields[0]]
            else:
                group_keys = []
                for group_field in group_fields:
                    if isinstance(item[group_field], int):
                        group_keys.append(str(item[group_field]))
                    if isinstance(item[group_field], unicode):
                        group_keys.append(item[group_field].encode('utf-8'))
                return '__'.join(group_keys)

        for item in iter(first_queryset):
            group_key = get_group_key(item, group_fields)
            result_dict[group_key] = set_item_keys(item.copy(), group_fields)

        for queryset in next_querysets:
            for item in iter(queryset):
                group_key = get_group_key(item, group_fields)

                # new key not previously seen
                if group_key not in result_dict:
                    result_dict[group_key] = set_item_keys(item.copy(), group_fields)
                else:
                    result_dict[group_key] = merge([result_dict[group_key], item.copy()])

        return list(result_dict.values())

    result = merge_results(aggregation_querysets, group_fields)

    return result

def serialize_foreign_keys(result, selected_groupings, request):
    """
    Re-use serializers to show full info of the grouped by items.

    Not all group by keys are serialized, this is based upon the value at _groupings.serializer
    """

    for grouping in selected_groupings:
        """
        Mutate result object for each item in the result[] array
        """
        serializer = grouping.serializer
        serializer_fields = grouping.serializer_fields

        result = grouping.serialize_results(result, request)

    return result

def apply_ordering(result, orderings):
    """
    orders a list by a key

    parameters
    result    - list of results
    orderings - list of order keys
    """

    if len(orderings):
        orderings = reversed(orderings)
        reverse = False

        for order in orderings:
            field = order

            if field[0] == '-':
                reverse = True
                field = field[1:]

            result = sorted(result, key=itemgetter(field), reverse=reverse)

    return result

def apply_group_filters(queryset, selected_groupings, params):
    """
    Filters that are applied only to filter direct visible results as returned
    by the GROUP_BY clause.
    """

    groupings = filter(
        lambda x: x.query_param in params,
        selected_groupings
    )

    for group in groupings:
        main_field = group.fields[0] # the one giving the relation from activity to id of item
        value = params[group.query_param]

        # TODO: We assume here all item filters are IN filters - 2016-03-07
        if isinstance(main_field, str):
            queryset = queryset.filter(**{"{}__in".format(main_field): value.split(',')})

    return queryset

def aggregate(queryset, request, selected_groupings, selected_aggregations, selected_orderings):
    """
        A view can call this function
    """
    # remove any existing ordering
    queryset = queryset.order_by()
    params = request.query_params

    # TODO: just throw exceptions here and catch in view - 2016-04-08

    if not len(selected_groupings):
        raise ValueError("Invalid value {} for mandatory field 'group_by'".format(params.get('group_by')))
    elif not len(selected_aggregations):
        raise ValueError("Invalid value {} for mandatory field 'aggregations'".format(params.get('aggregations')))

    # filters that reduce the amount of "items" returned in the group_by
    # These filters must be applied directly instead of through "activity id" IN filters
    queryset = apply_group_filters(queryset, selected_groupings, params)

    # from here, queryset is a list
    result = apply_annotations(queryset, selected_groupings, selected_aggregations, params)

    # TODO: is this correct? - 2016-04-07
    count = len(result)

    # TODO: Can we let db do the ordering? - 2016-04-07
    result = apply_ordering(result, selected_orderings)
    result = serialize_foreign_keys(result, selected_groupings, request)

    return {
        'count': count,
        'results': result
    }

def intersection(list1, list2):
    """
    Return the intersection of two lists
    """
    return list(set(list1) & set(list2))

def flatten(l):
    """
    flatten a list of lists to a single list
    """
    return reduce(lambda x, y: x+y, l)

def merge(l):
    """
    merge a list dictionaries to one dictionary
    """
    if len(l) is 1:
        return l[0]

    result = {}
    for d in l:
        result.update(d)

    return result

<commit_msg>extra group by with filter on same field bugfix<commit_after>
from operator import itemgetter

from django.db.models import Count
from django.db.models import Sum
from django.db.models import Q, F
from django.db.models.functions import Coalesce
from django.db import connection

def apply_annotations(queryset, selected_groupings, selected_aggregations, query_params):
    """
    Builds and performs the query, when multiple aggregations were requested it joins the results
    """
    result_dict = None

    #
    # Apply group_by fields and renames
    #

    group_fields = flatten([ grouping.get_fields() for grouping in selected_groupings ])
    rename_annotations = merge([ grouping.get_renamed_fields() for grouping in selected_groupings ])
    group_extras = merge([ grouping.extra for grouping in selected_groupings if grouping.extra is not None])

    eliminate_nulls = {"{}__isnull".format(grouping): False for grouping in group_fields}

    queryset = queryset \
        .extra(**group_extras) \
        .annotate(**rename_annotations)

    if len(group_extras) is 0:
        queryset = queryset \
            .filter(**eliminate_nulls)

    # preparation for aggregation look
    main_group_key = group_fields[0]
    rest_group_keys = group_fields[1:]

    def get_aggregation_queryset(queryset, group_fields, aggregation):

        # TODO: Should queryset be copied here? - 2016-04-07
        next_result = queryset.all()

        # apply any extra aggregation filters if specified
        next_result = aggregation.apply_extra_filter(next_result)

        # apply group_by values() call
        next_result = next_result.values(*group_fields)

        # apply the aggregation annotation
        next_result = aggregation.apply_annotation(next_result, query_params, selected_groupings)
        # print str(next_result.query)
        return next_result

    aggregation_querysets = [
        get_aggregation_queryset(queryset, group_fields, aggregation)
            for aggregation in selected_aggregations
    ]

    def merge_results(querysets, group_fields):
        """
        Execute the querysets and merge the results into one list of dictionaries
        This method keeps ordering of keys in order of execution of the aggregations
        """

        if len(querysets) is 1:
            return querysets[0]

        first_queryset = querysets[0]
        next_querysets = querysets[1:]

        result_dict = {}

        def set_item_keys(item, group_fields):

            for aggregation in selected_aggregations:
                if not aggregation.field in item:
                    item[aggregation.field] = 0
            return item

        def get_group_key(item, group_fields):
            """
            Generate a unique group key from fields that are grouped by
            """
            if len(group_fields) is 1:
                return item[group_fields[0]]
            else:
                group_keys = []
                for group_field in group_fields:
                    if isinstance(item[group_field], int):
                        group_keys.append(str(item[group_field]))
                    if isinstance(item[group_field], unicode):
                        group_keys.append(item[group_field].encode('utf-8'))
                return '__'.join(group_keys)

        for item in iter(first_queryset):
            group_key = get_group_key(item, group_fields)
            result_dict[group_key] = set_item_keys(item.copy(), group_fields)

        for queryset in next_querysets:
            for item in iter(queryset):
                group_key = get_group_key(item, group_fields)

                # new key not previously seen
                if group_key not in result_dict:
                    result_dict[group_key] = set_item_keys(item.copy(), group_fields)
                else:
                    result_dict[group_key] = merge([result_dict[group_key], item.copy()])

        return list(result_dict.values())

    result = merge_results(aggregation_querysets, group_fields)

    return result

def serialize_foreign_keys(result, selected_groupings, request):
    """
    Re-use serializers to show full info of the grouped by items.

    Not all group by keys are serialized, this is based upon the value at _groupings.serializer
    """

    for grouping in selected_groupings:
        """
        Mutate result object for each item in the result[] array
        """
        serializer = grouping.serializer
        serializer_fields = grouping.serializer_fields

        result = grouping.serialize_results(result, request)

    return result

def apply_ordering(result, orderings):
    """
    orders a list by a key

    parameters
    result    - list of results
    orderings - list of order keys
    """

    if len(orderings):
        orderings = reversed(orderings)
        reverse = False

        for order in orderings:
            field = order

            if field[0] == '-':
                reverse = True
                field = field[1:]

            result = sorted(result, key=itemgetter(field), reverse=reverse)

    return result

def apply_group_filters(queryset, selected_groupings, params):
    """
    Filters that are applied only to filter direct visible results as returned
    by the GROUP_BY clause.
    """

    groupings = filter(
        lambda x: x.query_param in params,
        selected_groupings
    )

    for group in groupings:
        if hasattr(group, 'extra'):
            continue

        main_field = group.fields[0] # the one giving the relation from activity to id of item
        value = params[group.query_param]

        # TODO: We assume here all item filters are IN filters - 2016-03-07
        if isinstance(main_field, str):
            queryset = queryset.filter(**{"{}__in".format(main_field): value.split(',')})

    return queryset

def aggregate(queryset, request, selected_groupings, selected_aggregations, selected_orderings):
    """
        A view can call this function
    """
    # remove any existing ordering
    queryset = queryset.order_by()
    params = request.query_params

    # TODO: just throw exceptions here and catch in view - 2016-04-08

    if not len(selected_groupings):
        raise ValueError("Invalid value {} for mandatory field 'group_by'".format(params.get('group_by')))
    elif not len(selected_aggregations):
        raise ValueError("Invalid value {} for mandatory field 'aggregations'".format(params.get('aggregations')))

    # filters that reduce the amount of "items" returned in the group_by
    # These filters must be applied directly instead of through "activity id" IN filters
    queryset = apply_group_filters(queryset, selected_groupings, params)

    # from here, queryset is a list
    result = apply_annotations(queryset, selected_groupings, selected_aggregations, params)

    # TODO: is this correct? - 2016-04-07
    count = len(result)

    # TODO: Can we let db do the ordering? - 2016-04-07
    result = apply_ordering(result, selected_orderings)
    result = serialize_foreign_keys(result, selected_groupings, request)

    return {
        'count': count,
        'results': result
    }

def intersection(list1, list2):
    """
    Return the intersection of two lists
    """
    return list(set(list1) & set(list2))

def flatten(l):
    """
    flatten a list of lists to a single list
    """
    return reduce(lambda x, y: x+y, l)

def merge(l):
    """
    merge a list dictionaries to one dictionary
    """
    if len(l) is 1:
        return l[0]

    result = {}
    for d in l:
        result.update(d)

    return result

<|endoftext|>


<commit_before># python3
# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Simple JAX actors."""

from typing import Callable, Generic, Optional, Tuple, TypeVar, Union

from acme import adders
from acme import core
from acme import types
from acme.agents.jax import actor_core
from acme.jax import networks as network_lib
from acme.jax import utils
from acme.jax import variable_utils
import dm_env
import jax

# Useful type aliases.
RecurrentState = TypeVar('RecurrentState')

# Signatures for functions that sample from parameterised stochastic policies.
FeedForwardPolicy = Callable[
    [network_lib.Params, network_lib.PRNGKey, network_lib.Observation],
    Union[network_lib.Action, Tuple[network_lib.Action, types.NestedArray]]]
RecurrentPolicy = Callable[[
    network_lib.Params, network_lib.PRNGKey, network_lib
    .Observation, RecurrentState
], Tuple[Union[network_lib.Action, Tuple[network_lib.Action,
                                         types.NestedArray]], RecurrentState]]

class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
  """A generic actor implemented on top of ActorCore.

  An actor based on a policy which takes observations and outputs actions. It
  also adds experiences to replay and updates the actor weights from the policy
  on the learner.
  """

  def __init__(
      self,
      actor: actor_core.ActorCore[actor_core.State, actor_core.Extras],
      random_key: network_lib.PRNGKey,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a feed forward actor.

    Args:
      actor: actor core.
      random_key: Random key.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key

    self._init = jax.jit(actor.init)
    self._policy = jax.jit(actor.select_action, backend=backend)
    self._get_extras = actor.get_extras

    self._adder = adder
    self._state = None
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> types.NestedArray:
    action, self._state = self._policy(self._client.params,
                                       observation, self._state)
    return utils.to_numpy(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    self._random_key, key = jax.random.split(self._random_key)
    self._state = self._init(key)
    if self._adder:
      self._adder.add_first(timestep)

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add(action, next_timestep,
                      extras=self._get_extras(self._state))

  def update(self, wait: bool = False):
    self._client.update(wait)

# TODO(raveman): Migrate all users of FeedForwardActor to GenericActor and
# remove this class.
class FeedForwardActor(core.Actor):
  """A simple feed-forward actor implemented in JAX.

  An actor based on a policy which takes observations and outputs actions. It
  also adds experiences to replay and updates the actor weights from the policy
  on the learner.
  """

  def __init__(
      self,
      policy: FeedForwardPolicy,
      random_key: network_lib.PRNGKey,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      has_extras: bool = False,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a feed forward actor.

    Args:
      policy: A policy network taking observation and returning an action, if
        `has_extras=False`, and returning an (action, extras) tuple if
        `has_extras=True`.
      random_key: Random key.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to.
      has_extras: Flag indicating whether the policy returns extra
        information (e.g. q-values) in addition to an action.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key
    self._has_extras = has_extras
    self._extras: types.NestedArray = ()

    # Adding batch dimension inside jit is much more efficient than outside.
    def batched_policy(
        params: network_lib.Params, key: network_lib.PRNGKey,
        observation: network_lib.Observation
    ) -> Tuple[Union[network_lib.Action, Tuple[
        network_lib.Action, types.NestedArray]], network_lib.PRNGKey]:
      # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
      key, key2 = jax.random.split(key)
      observation = utils.add_batch_dim(observation)
      output = policy(params, key2, observation)
      return utils.squeeze_batch_dim(output), key
    self._policy = jax.jit(batched_policy, backend=backend)

    self._adder = adder
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> types.NestedArray:
    result, self._random_key = self._policy(self._client.params,
                                            self._random_key, observation)
    if self._has_extras:
      action, self._extras = result
    else:
      action = result
    return utils.to_numpy(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add_first(timestep)

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add(action, next_timestep, extras=self._extras)

  def update(self, wait: bool = False):
    self._client.update(wait)

# TODO(raveman): Migrate all users of RecurrentActor to GenericActor and
# remove this class.
class RecurrentActor(core.Actor):
  """A recurrent actor in JAX.

  An actor based on a recurrent policy which takes observations and outputs
  actions, and keeps track of the recurrent state inside. It also adds
  experiences to replay and updates the actor weights from the policy on the
  learner.
  """

  def __init__(
      self,
      recurrent_policy: RecurrentPolicy,
      random_key: network_lib.PRNGKey,
      initial_core_state: RecurrentState,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      has_extras: bool = False,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a recurrent actor.

    Args:
      recurrent_policy: A recurrent policy network taking observation and state
        and returning an (action, state) tuple, if `has_extras=False`, and
        returning an ((action, extras), state) tuple if `has_extras=True`. In
        the latter case, `extras` must be a tuple.
      random_key: Random key.
      initial_core_state: Initial state of the recurrent policy.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to. The `extras` of the adder holds the
        state of the recurrent policy. If `has_extras=True` then the `extras`
        part returned from the recurrent policy is appended to the state before
        added to the adder.
      has_extras: Flag indicating whether the recurrent policy returns extra
        information (e.g. q-values) in addition to an action.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key
    self._has_extras = has_extras
    self._extras: types.NestedArray = ()

    # Adding batch dimension inside jit is much more efficient than outside.
    def batched_recurrent_policy(
        params: network_lib.Params, key: network_lib.PRNGKey,
        observation: network_lib.Observation, core_state: RecurrentState
    ) -> Tuple[Union[network_lib.Action, Tuple[
        network_lib.Action, types.NestedArray]], RecurrentState,
               network_lib.PRNGKey]:
      # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
      observation = utils.add_batch_dim(observation)
      key, key2 = jax.random.split(key)
      output, new_state = recurrent_policy(params, key2, observation,
                                           core_state)
      return output, new_state, key
    self._recurrent_policy = jax.jit(batched_recurrent_policy, backend=backend)

    self._initial_state = self._prev_state = self._state = initial_core_state
    self._adder = adder
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> network_lib.Action:
    result, new_state, self._random_key = self._recurrent_policy(
        self._client.params,
        key=self._random_key,
        observation=observation,
        core_state=self._state)
    self._prev_state = self._state  # Keep previous state to save in replay.
    self._state = new_state  # Keep new state for next policy call.

    if self._has_extras:
      action, extras = result
      self._extras = utils.to_numpy_squeeze(extras)   # Keep to save in replay.
    else:
      action = result
    return utils.to_numpy_squeeze(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add_first(timestep)
    # Re-initialize state at beginning of new episode.
    self._state = self._initial_state

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      numpy_state = utils.to_numpy_squeeze(self._prev_state)
      self._adder.add(
          action, next_timestep, extras=(numpy_state,) + self._extras)

  def update(self, wait: bool = False):
    self._client.update(wait)
<commit_msg>Change how the RecurrentActor handles extras in its observe method.<commit_after># python3
# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Simple JAX actors."""

from typing import Callable, Generic, Optional, Tuple, TypeVar, Union

from acme import adders
from acme import core
from acme import types
from acme.agents.jax import actor_core
from acme.jax import networks as network_lib
from acme.jax import utils
from acme.jax import variable_utils
import dm_env
import jax

# Useful type aliases.
RecurrentState = TypeVar('RecurrentState')

# Signatures for functions that sample from parameterised stochastic policies.
FeedForwardPolicy = Callable[
    [network_lib.Params, network_lib.PRNGKey, network_lib.Observation],
    Union[network_lib.Action, Tuple[network_lib.Action, types.NestedArray]]]
RecurrentPolicy = Callable[[
    network_lib.Params, network_lib.PRNGKey, network_lib
    .Observation, RecurrentState
], Tuple[Union[network_lib.Action, Tuple[network_lib.Action,
                                         types.NestedArray]], RecurrentState]]

class GenericActor(core.Actor, Generic[actor_core.State, actor_core.Extras]):
  """A generic actor implemented on top of ActorCore.

  An actor based on a policy which takes observations and outputs actions. It
  also adds experiences to replay and updates the actor weights from the policy
  on the learner.
  """

  def __init__(
      self,
      actor: actor_core.ActorCore[actor_core.State, actor_core.Extras],
      random_key: network_lib.PRNGKey,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a feed forward actor.

    Args:
      actor: actor core.
      random_key: Random key.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key

    self._init = jax.jit(actor.init)
    self._policy = jax.jit(actor.select_action, backend=backend)
    self._get_extras = actor.get_extras

    self._adder = adder
    self._state = None
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> types.NestedArray:
    action, self._state = self._policy(self._client.params, observation,
                                       self._state)
    return utils.to_numpy(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    self._random_key, key = jax.random.split(self._random_key)
    self._state = self._init(key)
    if self._adder:
      self._adder.add_first(timestep)

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add(
          action, next_timestep, extras=self._get_extras(self._state))

  def update(self, wait: bool = False):
    self._client.update(wait)

# TODO(raveman): Migrate all users of FeedForwardActor to GenericActor and
# remove this class.
class FeedForwardActor(core.Actor):
  """A simple feed-forward actor implemented in JAX.

  An actor based on a policy which takes observations and outputs actions. It
  also adds experiences to replay and updates the actor weights from the policy
  on the learner.
  """

  def __init__(
      self,
      policy: FeedForwardPolicy,
      random_key: network_lib.PRNGKey,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      has_extras: bool = False,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a feed forward actor.

    Args:
      policy: A policy network taking observation and returning an action, if
        `has_extras=False`, and returning an (action, extras) tuple if
        `has_extras=True`.
      random_key: Random key.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to.
      has_extras: Flag indicating whether the policy returns extra information
        (e.g. q-values) in addition to an action.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key
    self._has_extras = has_extras
    self._extras: types.NestedArray = ()

    # Adding batch dimension inside jit is much more efficient than outside.
    def batched_policy(
        params: network_lib.Params, key: network_lib.PRNGKey,
        observation: network_lib.Observation
    ) -> Tuple[Union[network_lib.Action, Tuple[
        network_lib.Action, types.NestedArray]], network_lib.PRNGKey]:
      # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
      key, key2 = jax.random.split(key)
      observation = utils.add_batch_dim(observation)
      output = policy(params, key2, observation)
      return utils.squeeze_batch_dim(output), key

    self._policy = jax.jit(batched_policy, backend=backend)

    self._adder = adder
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> types.NestedArray:
    result, self._random_key = self._policy(self._client.params,
                                            self._random_key, observation)
    if self._has_extras:
      action, self._extras = result
    else:
      action = result
    return utils.to_numpy(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add_first(timestep)

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add(action, next_timestep, extras=self._extras)

  def update(self, wait: bool = False):
    self._client.update(wait)

# TODO(raveman): Migrate all users of RecurrentActor to GenericActor and
# remove this class.
class RecurrentActor(core.Actor):
  """A recurrent actor in JAX.

  An actor based on a recurrent policy which takes observations and outputs
  actions, and keeps track of the recurrent state inside. It also adds
  experiences to replay and updates the actor weights from the policy on the
  learner.
  """

  def __init__(
      self,
      recurrent_policy: RecurrentPolicy,
      random_key: network_lib.PRNGKey,
      initial_core_state: RecurrentState,
      variable_client: variable_utils.VariableClient,
      adder: Optional[adders.Adder] = None,
      has_extras: bool = False,
      backend: Optional[str] = 'cpu',
  ):
    """Initializes a recurrent actor.

    Args:
      recurrent_policy: A recurrent policy network taking observation and state
        and returning an (action, state) tuple, if `has_extras=False`, and
        returning an ((action, extras), state) tuple if `has_extras=True`. In
        the latter case, `extras` must be a tuple.
      random_key: Random key.
      initial_core_state: Initial state of the recurrent policy.
      variable_client: The variable client to get policy parameters from.
      adder: An adder to add experiences to. The `extras` of the adder holds the
        state of the recurrent policy. If `has_extras=True` then the `extras`
        part returned from the recurrent policy is appended to the state before
        added to the adder.
      has_extras: Flag indicating whether the recurrent policy returns extra
        information (e.g. q-values) in addition to an action.
      backend: Which backend to use for running the policy.
    """
    self._random_key = random_key
    self._has_extras = has_extras
    self._extras: types.NestedArray = ()

    # Adding batch dimension inside jit is much more efficient than outside.
    def batched_recurrent_policy(
        params: network_lib.Params, key: network_lib.PRNGKey,
        observation: network_lib.Observation, core_state: RecurrentState
    ) -> Tuple[Union[network_lib.Action, Tuple[
        network_lib.Action, types.NestedArray]], RecurrentState,
               network_lib.PRNGKey]:
      # TODO(b/161332815): Make JAX Actor work with batched or unbatched inputs.
      observation = utils.add_batch_dim(observation)
      key, key2 = jax.random.split(key)
      output, new_state = recurrent_policy(params, key2, observation,
                                           core_state)
      return output, new_state, key

    self._recurrent_policy = jax.jit(batched_recurrent_policy, backend=backend)

    self._initial_state = self._prev_state = self._state = initial_core_state
    self._adder = adder
    self._client = variable_client

  def select_action(self,
                    observation: network_lib.Observation) -> network_lib.Action:
    result, new_state, self._random_key = self._recurrent_policy(
        self._client.params,
        key=self._random_key,
        observation=observation,
        core_state=self._state)
    self._prev_state = self._state  # Keep previous state to save in replay.
    self._state = new_state  # Keep new state for next policy call.

    if self._has_extras:
      action, extras = result
      self._extras = utils.to_numpy_squeeze(extras)  # Keep to save in replay.
    else:
      action = result
    return utils.to_numpy_squeeze(action)

  def observe_first(self, timestep: dm_env.TimeStep):
    if self._adder:
      self._adder.add_first(timestep)
    # Re-initialize state at beginning of new episode.
    self._state = self._initial_state

  def observe(self, action: network_lib.Action, next_timestep: dm_env.TimeStep):
    if self._adder:
      # Convert state to numpy array and pack it in a dictionary.
      extras = {'core_state': utils.to_numpy_squeeze(self._prev_state)}

      # Add core state to the extras dict or tuple, accordingly.
      if self._has_extras:
        if 'core_state' in self._extras:
          raise ValueError(
              'The policy network\'s extras dict already has a `core_state` '
              'but the actor is attempting to add its own `core_state`.')
        extras.update(self._extras)

      # Add the transition to the adder.
      self._adder.add(action, next_timestep, extras=extras)

  def update(self, wait: bool = False):
    self._client.update(wait)
<|endoftext|>


<commit_before>import inspect
import time
from datetime import timedelta

from django.core.management.base import BaseCommand, CommandError

from corehq.elastic import get_es_export, ES_META
from pillowtop.es_utils import initialize_index, set_index_reindex_settings

USAGE = """Reindex data from one ES index into another ES index using Elasticsearch reindex API
    https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-reindex.html#reindex-from-remote.
"""

class Command(BaseCommand):
    help = USAGE

    def add_arguments(self, parser):
        parser.add_argument(
            "source_index",
            help="Index to be used as the source",
        )
        parser.add_argument(
            "target_index", choices=ES_META,
            help="Index to be used as the target"
        )
        parser.add_argument(
            "--print_index_size",
            action="store_true",
            default=False,
            help="Print number of docs in source and target index"
        )

    def handle(self, source_index, target_index, **options):
        es = get_es_export()
        if not es.indices.exists(source_index):
            raise CommandError(f"Source index does not exist: '{source_index}'")

        target_index_info = ES_META[target_index]
        target_index = target_index_info.index
        _initialize_target(es, target_index_info)

        source_count = _get_doc_count(es, source_index)
        target_count = _get_doc_count(es, target_index)
        print("Number of docs in source index", source_count)
        print("Number of docs in target index", target_count)

        if options.get("print_index_size"):
            return

        elif source_count == target_count:
            print("Doc counts are same, nothing left to reindex. Exiting!")
            return

        print(f"Starting reindex for index from '{source_index}' to '{target_index}'")
        reindex_and_poll_progress(es, source_index, target_index)

        print(inspect.cleandoc(f"""
            If you would like to check, you may run this
            command again with --print_index_size to make sure
            doc counts match across both indices.

            Once you are satisfied that the reindex is complete you should run the following:

            ./manage.py initialize_es_indices --index {target_index} --set-for-usage
        """))

def _initialize_target(es, target_index_info):
    if es.indices.exists(target_index_info.index):
        print(f"Target index '{target_index_info.index}' already exists. Skipping initialization.")
        return

    print("Creating target index")
    initialize_index(es, target_index_info)
    set_index_reindex_settings(es, target_index_info.index)

    print("Setting number of replicas to 0")
    es.indices.put_settings({
        "index.number_of_replicas": 0
    }, index=target_index_info.index)

def reindex_and_poll_progress(es, source_index, target_index):
    reindex_query = {
        "source": {
            "index": source_index,
        },
        "dest": {
            "index": target_index
        },
        "conflicts": "proceed"
    }

    result = es.reindex(reindex_query, wait_for_completion=False, request_timeout=300)
    task_id = result["task"]
    node_id = task_id.split(':')[0]
    node_name = es.nodes.info(node_id, metric="name")["nodes"][node_id]["name"]
    print(f"Task with ID '{task_id}' running on '{node_name}'")
    while True:
        result = es.tasks.list(task_id=task_id, detailed=True)
        if not result["nodes"]:
            node_failure = result["node_failures"][0]
            error = node_failure["caused_by"]["type"]
            if error == "resource_not_found_exception":
                print("Reindex task complete or failed. Please check ES logs for final output.")
                return
            else:
                raise CommandError(f"Fetching task failed: {node_failure}")

        task_details = result["nodes"][node_id]["tasks"][task_id]
        status = task_details["status"]
        total = status["total"]
        created, updated, deleted = status["created"], status["updated"], status["deleted"]
        progress = created + updated + deleted
        progress_percent = progress / total * 100

        running_time_nanos = task_details["running_time_in_nanos"]
        run_time = timedelta(microseconds=running_time_nanos / 1000)

        print(f"Progress {progress_percent:.2f}% ({progress} / {total}). Duration: {run_time}")

        time.sleep(5)

def _get_doc_count(es, index):
    es.indices.refresh(index)
    return es.indices.stats(index=index)['indices'][index]['primaries']['docs']['count']
<commit_msg>option to monitor existing reindex progress<commit_after>import inspect
import time
from datetime import timedelta

from django.core.management.base import BaseCommand, CommandError

from corehq.elastic import get_es_export, ES_META
from pillowtop.es_utils import initialize_index, set_index_reindex_settings

USAGE = """Reindex data from one ES index into another ES index using Elasticsearch reindex API
    https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-reindex.html#reindex-from-remote.
"""

class Command(BaseCommand):
    help = USAGE

    def add_arguments(self, parser):
        parser.add_argument(
            "source_index",
            help="Index to be used as the source",
        )
        parser.add_argument(
            "target_index", choices=ES_META,
            help="Index to be used as the target"
        )
        parser.add_argument(
            "--print_index_size",
            action="store_true",
            default=False,
            help="Print number of docs in source and target index"
        )
        parser.add_argument(
            "--monitor-task-id",
            help="Pass in the task ID of an existing reindex to monitor it's progress. This will not kick off"
                 "a new reindex."
        )

    def handle(self, source_index, target_index, **options):
        es = get_es_export()
        if not es.indices.exists(source_index):
            raise CommandError(f"Source index does not exist: '{source_index}'")

        target_index_info = ES_META[target_index]
        target_index = target_index_info.index
        _initialize_target(es, target_index_info)

        source_count = _get_doc_count(es, source_index)
        target_count = _get_doc_count(es, target_index)
        print("Number of docs in source index", source_count)
        print("Number of docs in target index", target_count)

        if options.get("print_index_size"):
            return

        elif source_count == target_count:
            print("Doc counts are same, nothing left to reindex. Exiting!")
            return

        if not options["monitor_task_id"]:
            print(f"Starting reindex for index from '{source_index}' to '{target_index}'")
            task_id = start_reindex(es, source_index, target_index)
        else:
            task_id = options["monitor_task_id"]
        check_task_progress(es, task_id)

        print(inspect.cleandoc(f"""
            If you would like to check, you may run this
            command again with --print_index_size to make sure
            doc counts match across both indices.

            Once you are satisfied that the reindex is complete you should run the following:

            ./manage.py initialize_es_indices --index {target_index} --set-for-usage
        """))

def _initialize_target(es, target_index_info):
    if es.indices.exists(target_index_info.index):
        print(f"Target index '{target_index_info.index}' already exists. Skipping initialization.")
        return

    print("Creating target index")
    initialize_index(es, target_index_info)
    set_index_reindex_settings(es, target_index_info.index)

    print("Setting number of replicas to 0")
    es.indices.put_settings({
        "index.number_of_replicas": 0
    }, index=target_index_info.index)

def start_reindex(es, source_index, target_index):
    reindex_query = {
        "source": {
            "index": source_index,
        },
        "dest": {
            "index": target_index
        },
        "conflicts": "proceed"
    }

    result = es.reindex(reindex_query, wait_for_completion=False, request_timeout=300)
    task_id = result["task"]
    return task_id

def check_task_progress(es, task_id):
    node_id = task_id.split(':')[0]
    node_name = es.nodes.info(node_id, metric="name")["nodes"][node_id]["name"]
    print(f"Task with ID '{task_id}' running on '{node_name}'")
    while True:
        result = es.tasks.list(task_id=task_id, detailed=True)
        if not result["nodes"]:
            node_failure = result["node_failures"][0]
            error = node_failure["caused_by"]["type"]
            if error == "resource_not_found_exception":
                print("Reindex task complete or failed. Please check ES logs for final output.")
                return
            else:
                raise CommandError(f"Fetching task failed: {node_failure}")

        task_details = result["nodes"][node_id]["tasks"][task_id]
        status = task_details["status"]
        total = status["total"]
        created, updated, deleted = status["created"], status["updated"], status["deleted"]
        progress = created + updated + deleted
        progress_percent = progress / total * 100

        running_time_nanos = task_details["running_time_in_nanos"]
        run_time = timedelta(microseconds=running_time_nanos / 1000)

        print(f"Progress {progress_percent:.2f}% ({progress} / {total}). Duration: {run_time}")

        time.sleep(5)

def _get_doc_count(es, index):
    es.indices.refresh(index)
    return es.indices.stats(index=index)['indices'][index]['primaries']['docs']['count']
<|endoftext|>