<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Interpreter version: python 2.7
#
# Imports =====================================================================
import pytest
import dhtmlparser

from harvester.scrappers import cpress_cz

# Variables ===================================================================

# Functions & objects =========================================================
def test_parse_alt_title():
    dom = dhtmlparser.parseString(
        """

        """
    )

    alt_title = cpress_cz._parse_alt_title(dom)

    assert alt_title == "Záhadná jízda králů"

def test_parse_alt_title_not_found():
    dom = dhtmlparser.parseString(
        """

        """
    )

    with pytest.raises(UserWarning):
        cpress_cz._parse_alt_title(dom)

def test_parse_alt_title_param_not_found():
    dom = dhtmlparser.parseString(
        """

        """
    )

    with pytest.raises(UserWarning):
        cpress_cz._parse_alt_title(dom)

def test_parse_alt_url():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    url = cpress_cz._parse_alt_url(dom)

    assert url
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_alt_url_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

        """
    )

    url = cpress_cz._parse_alt_url(dom)

    assert url is None

def test_parse_title_url():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_title_url_tag_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_title_url_url_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url is None

def test_parse_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors
    assert authors[0].name == "Jiří Jilík"
    assert authors[0].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/jiri-jilik/")

def test_parse_authors_no_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors == []

def test_parse_authors_no_authors_tag():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors == []

def test_parse_authors_multiple_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Leoš Kopecký, Roswitha Kammerl

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors
    assert len(authors) == 2

    assert authors[0].name == "Leoš Kopecký"
    assert authors[0].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/leos-kopecky/")
    assert authors[1].name == "Roswitha Kammerl"
    assert authors[1].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/roswitha-kammerl/")<commit_msg>#15, #11; Added tests of _parse_price().<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# Interpreter version: python 2.7
#
# Imports =====================================================================
import pytest
import dhtmlparser

from harvester.scrappers import cpress_cz

# Variables ===================================================================

# Functions & objects =========================================================
def test_parse_alt_title():
    dom = dhtmlparser.parseString(
        """

        """
    )

    alt_title = cpress_cz._parse_alt_title(dom)

    assert alt_title == "Záhadná jízda králů"

def test_parse_alt_title_not_found():
    dom = dhtmlparser.parseString(
        """

        """
    )

    with pytest.raises(UserWarning):
        cpress_cz._parse_alt_title(dom)

def test_parse_alt_title_param_not_found():
    dom = dhtmlparser.parseString(
        """

        """
    )

    with pytest.raises(UserWarning):
        cpress_cz._parse_alt_title(dom)

def test_parse_alt_url():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    url = cpress_cz._parse_alt_url(dom)

    assert url
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_alt_url_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

        """
    )

    url = cpress_cz._parse_alt_url(dom)

    assert url is None

def test_parse_title_url():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_title_url_tag_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url == cpress_cz.normalize_url(cpress_cz.BASE_URL, "zahadna-jizda-kralu/")

def test_parse_title_url_url_not_found():
    dom = dhtmlparser.parseString(
        """

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    title, url = cpress_cz._parse_title_url(dom)

    assert title == "Záhadná jízda králů"
    assert url is None

def test_parse_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Jiří Jilík

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors
    assert authors[0].name == "Jiří Jilík"
    assert authors[0].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/jiri-jilik/")

def test_parse_authors_no_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors == []

def test_parse_authors_no_authors_tag():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors == []

def test_parse_authors_multiple_authors():
    dom = dhtmlparser.parseString(
        """

        Záhadná jízda králů

Leoš Kopecký, Roswitha Kammerl

Nová kniha autora bestselleru Žítkovské čarování!

199 Kč

Detail knihy

        """
    )

    authors = cpress_cz._parse_authors(dom)

    assert authors
    assert len(authors) == 2

    assert authors[0].name == "Leoš Kopecký"
    assert authors[0].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/leos-kopecky/")
    assert authors[1].name == "Roswitha Kammerl"
    assert authors[1].URL == cpress_cz.normalize_url(cpress_cz.BASE_URL, "autori/autor/roswitha-kammerl/")

def test_parse_price():
    dom = dhtmlparser.parseString(
        """

((availability)): ((availability_available))

Doporučená cena: 299 Kč

        """
    )

    price = cpress_cz._parse_price(dom)

    assert price == "299 Kč"

def test_parse_price_not_found():
    dom = dhtmlparser.parseString(
        """

((availability)): ((availability_available))

        """
    )

    with pytest.raises(UserWarning):
        cpress_cz._parse_price(dom)
<|endoftext|>


<commit_before># Copyright 2017 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging

from google.appengine.api import oauth

from dashboard.common import datastore_hooks
from dashboard.common import utils

OAUTH_CLIENT_ID_WHITELIST = [
    # This oauth client id is from Pinpoint.
    '62121018386-aqdfougp0ddn93knqj6g79vvn42ajmrg.apps.googleusercontent.com',
    # This oauth client id is from the 'chromeperf' API console.
    '62121018386-h08uiaftreu4dr3c4alh3l7mogskvb7i.apps.googleusercontent.com',
    # This oauth client id is from chromiumdash-staging.
    '377415874083-slpqb5ur4h9sdfk8anlq4qct9imivnmt.apps.googleusercontent.com',
    # This oauth client id is from chromiumdash.
    '975044924533-p122oecs8h6eibv5j5a8fmj82b0ct0nk.apps.googleusercontent.com',
    # This oauth client id is used to upload histograms from the perf waterfall.
    '113172445342431053212',
    'chromeperf@webrtc-perf-test.google.com.iam.gserviceaccount.com',
    # This oauth client id is used to upload histograms when debugging Fuchsia
    # locally (e.g. in a cron-job).
    'catapult-uploader@fuchsia-infra.iam.gserviceaccount.com',
    # This oauth client id is used to upload histograms from Fuchsia dev
    # builders.
    'garnet-ci-builder-dev@fuchsia-infra.iam.gserviceaccount.com',
    # These client id's are used from Fuchsia CI builders.
    'garnet-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'global-integration-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'fuchsia-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'peridot-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'topaz-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'vendor-google-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    # This oauth client id used to upload histograms from cronet bots.
    '113172445342431053212',
    # Used by luci builders to upload perf data.
    'chrome-ci-builder@chops-service-accounts.iam.gserviceaccount.com',
    # Used by luci builders to upload perf data.
    'chrome-try-builder@chops-service-accounts.iam.gserviceaccount.com',
    # These oauth client ids are used to upload Android performance metrics.
    '528014426327-fptk0tpfi4orpcol559k77v7bi9onpq5.apps.googleusercontent.com',
    ('android-metrics-dashboard'
     '@android-metrics-dashboard.iam.gserviceaccount.com'),
    '107857144893180953937',
    # This oauth client id is used by all LUCI binaries. In particular, it will
    # allow accessing the APIs by authorized users that generate tokens via
    # luci-auth command.
    '446450136466-2hr92jrq8e6i4tnsa56b52vacp7t3936.apps.googleusercontent.com',
    # This oauth client id is used to upload ChromeOS performance metrics.
    'chromeperf-uploader@chromeos-bot.iam.gserviceaccount.com',
]

class ApiAuthException(Exception):
  pass

class OAuthError(ApiAuthException):
  def __init__(self):
    super(OAuthError, self).__init__('User authentication error')

class NotLoggedInError(ApiAuthException):
  def __init__(self):
    super(NotLoggedInError, self).__init__('User not authenticated')

class InternalOnlyError(ApiAuthException):
  def __init__(self):
    super(InternalOnlyError, self).__init__('User does not have access')

def Authorize():
  try:
    email = utils.GetEmail()
  except oauth.OAuthRequestError:
    raise OAuthError

  if not email:
    raise NotLoggedInError

  try:
    if not email.endswith('.gserviceaccount.com'):
      # For non-service account, need to verify that the OAuth client ID
      # is in our whitelist.
      client_id = oauth.get_client_id(utils.OAUTH_SCOPES)
      if client_id not in OAUTH_CLIENT_ID_WHITELIST:
        logging.error('OAuth client id %s for user %s not in whitelist',
                      client_id, email)
        email = None
        raise OAuthError
  except oauth.OAuthRequestError:
    # Transient errors when checking the token result should result in HTTP 500,
    # so catch oauth.OAuthRequestError here, not oauth.Error (which would catch
    # both fatal and transient errors).
    raise OAuthError

  logging.info('OAuth user logged in as: %s', email)
  if utils.IsInternalUser():
    datastore_hooks.SetPrivilegedRequest()
<commit_msg>Whitelist V8 CI task account for uploading isolate hashes to Pinpoint<commit_after># Copyright 2017 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging

from google.appengine.api import oauth

from dashboard.common import datastore_hooks
from dashboard.common import utils

OAUTH_CLIENT_ID_WHITELIST = [
    # This oauth client id is from Pinpoint.
    '62121018386-aqdfougp0ddn93knqj6g79vvn42ajmrg.apps.googleusercontent.com',
    # This oauth client id is from the 'chromeperf' API console.
    '62121018386-h08uiaftreu4dr3c4alh3l7mogskvb7i.apps.googleusercontent.com',
    # This oauth client id is from chromiumdash-staging.
    '377415874083-slpqb5ur4h9sdfk8anlq4qct9imivnmt.apps.googleusercontent.com',
    # This oauth client id is from chromiumdash.
    '975044924533-p122oecs8h6eibv5j5a8fmj82b0ct0nk.apps.googleusercontent.com',
    # This oauth client id is used to upload histograms from the perf waterfall.
    '113172445342431053212',
    'chromeperf@webrtc-perf-test.google.com.iam.gserviceaccount.com',
    # This oauth client id is used to upload histograms when debugging Fuchsia
    # locally (e.g. in a cron-job).
    'catapult-uploader@fuchsia-infra.iam.gserviceaccount.com',
    # This oauth client id is used to upload histograms from Fuchsia dev
    # builders.
    'garnet-ci-builder-dev@fuchsia-infra.iam.gserviceaccount.com',
    # These client id's are used from Fuchsia CI builders.
    'garnet-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'global-integration-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'fuchsia-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'peridot-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'topaz-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    'vendor-google-ci-builder@fuchsia-infra.iam.gserviceaccount.com',
    # This oauth client id used to upload histograms from cronet bots.
    '113172445342431053212',
    # Used by luci builders to upload perf data.
    'chrome-ci-builder@chops-service-accounts.iam.gserviceaccount.com',
    'v8-ci-builder@chops-service-accounts.iam.gserviceaccount.com',
    # Used by luci builders to upload perf data.
    'chrome-try-builder@chops-service-accounts.iam.gserviceaccount.com',
    'v8-try-builder@chops-service-accounts.iam.gserviceaccount.com',
    # These oauth client ids are used to upload Android performance metrics.
    '528014426327-fptk0tpfi4orpcol559k77v7bi9onpq5.apps.googleusercontent.com',
    ('android-metrics-dashboard'
     '@android-metrics-dashboard.iam.gserviceaccount.com'),
    '107857144893180953937',
    # This oauth client id is used by all LUCI binaries. In particular, it will
    # allow accessing the APIs by authorized users that generate tokens via
    # luci-auth command.
    '446450136466-2hr92jrq8e6i4tnsa56b52vacp7t3936.apps.googleusercontent.com',
    # This oauth client id is used to upload ChromeOS performance metrics.
    'chromeperf-uploader@chromeos-bot.iam.gserviceaccount.com',
]

class ApiAuthException(Exception):
  pass

class OAuthError(ApiAuthException):
  def __init__(self):
    super(OAuthError, self).__init__('User authentication error')

class NotLoggedInError(ApiAuthException):
  def __init__(self):
    super(NotLoggedInError, self).__init__('User not authenticated')

class InternalOnlyError(ApiAuthException):
  def __init__(self):
    super(InternalOnlyError, self).__init__('User does not have access')

def Authorize():
  try:
    email = utils.GetEmail()
  except oauth.OAuthRequestError:
    raise OAuthError

  if not email:
    raise NotLoggedInError

  try:
    if not email.endswith('.gserviceaccount.com'):
      # For non-service account, need to verify that the OAuth client ID
      # is in our whitelist.
      client_id = oauth.get_client_id(utils.OAUTH_SCOPES)
      if client_id not in OAUTH_CLIENT_ID_WHITELIST:
        logging.error('OAuth client id %s for user %s not in whitelist',
                      client_id, email)
        email = None
        raise OAuthError
  except oauth.OAuthRequestError:
    # Transient errors when checking the token result should result in HTTP 500,
    # so catch oauth.OAuthRequestError here, not oauth.Error (which would catch
    # both fatal and transient errors).
    raise OAuthError

  logging.info('OAuth user logged in as: %s', email)
  if utils.IsInternalUser():
    datastore_hooks.SetPrivilegedRequest()
<|endoftext|>


<commit_before># Copyright 2018 The Bazel Authors. 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.
load(
    "//kotlin/internal/jvm:compile.bzl",
    "export_only_providers",
    _kt_jvm_produce_jar_actions = "kt_jvm_produce_jar_actions",
)
load(
    "//kotlin/internal:defs.bzl",
    _KtCompilerPluginInfo = "KtCompilerPluginInfo",
    _KtJvmInfo = "KtJvmInfo",
)
load(
    "//kotlin/internal/utils:utils.bzl",
    _utils = "utils",
)
load("//third_party:jarjar.bzl", "jarjar_action")

def _make_providers(ctx, providers, transitive_files = depset(order = "default"), *additional_providers):
    return struct(
        kt = providers.kt,
        providers = [
            providers.java,
            providers.kt,
            DefaultInfo(
                files = depset([ctx.outputs.jar, ctx.outputs.jdeps]),
                runfiles = ctx.runfiles(
                    # explicitly include data files, otherwise they appear to be missing
                    files = ctx.files.data,
                    transitive_files = transitive_files,
                    # continue to use collect_default until proper transitive data collecting is
                    # implmented.
                    collect_default = True,
                ),
            ),
        ] + list(additional_providers),
    )

def _write_launcher_action(ctx, rjars, main_class, jvm_flags, args = "", wrapper_preamble = ""):
    """Macro that writes out a launcher script shell script.
      Args:
        rjars: All of the runtime jars required to launch this java target.
        main_class: the main class to launch.
        jvm_flags: The flags that should be passed to the jvm.
        args: Args that should be passed to the Binary.
    """
    classpath = ":".join(["${RUNPATH}%s" % (j.short_path) for j in rjars.to_list()])
    jvm_flags = " ".join([ctx.expand_location(f, ctx.attr.data) for f in jvm_flags])
    template = ctx.attr._java_stub_template.files.to_list()[0]
    javabin = "JAVABIN=" + str(ctx.attr._java_runtime[java_common.JavaRuntimeInfo].java_executable_exec_path)

    ctx.actions.expand_template(
        template = template,
        output = ctx.outputs.executable,
        substitutions = {
            "%classpath%": classpath,
            "%java_start_class%": main_class,
            "%javabin%": javabin,
            "%jvm_flags%": jvm_flags,
            "%set_jacoco_metadata%": "",
            "%set_jacoco_main_class%": "",
            "%set_jacoco_java_runfiles_root%": "",
            "%set_java_coverage_new_implementation%": "",
            "%workspace_prefix%": ctx.workspace_name + "/",
        },
        is_executable = True,
    )

def _is_source_jar_stub(jar):
    """Workaround for intellij plugin expecting a source jar"""
    return jar.path.endswith("third_party/empty.jar")

def _unify_jars(ctx):
    if bool(ctx.attr.jar):
        return struct(class_jar = ctx.file.jar, source_jar = ctx.file.srcjar, ijar = None)
    else:
        # Legacy handling.
        jars = []
        source_jars = []
        if (ctx.file.srcjar and not ("%s" % ctx.file.srcjar.path).endswith("third_party/empty.jar")):
            source_jars += [ctx.file.srcjar]

        # There must be a single jar jar and it can either be a filegroup or a JavaInfo.
        for jar in ctx.attr.jars:
            # If a JavaInfo is available it's because it was picked up from a `maven_jar` style attribute -- e.g.,
            # @com_google_guava_guava//jar. so the transitive_compile_jars or the transitive_runtime_jars should not be
            # visited -- descending into these results in ijars entering the graph.
            if JavaInfo in jar:
                jars += jar[JavaInfo].full_compile_jars.to_list()
                source_jars += jar[JavaInfo].transitive_source_jars.to_list()
            else:
                # this branch occurs when the attr was a filegroup.
                for file in jar.files.to_list():
                    if file.basename.endswith("-sources.jar"):
                        source_jars.append(file)
                    elif file.basename.endswith(".jar"):
                        jars.append(file)
                    else:
                        fail("a jar pointing to a filegroup must either end with -sources.jar or .jar")

        if len(jars) > 1:
            fail("Got more than one jar, this is an error create an issue: %s" % jars)
        if len(source_jars) > 1:
            fail("Got more than one source jar. " +
                 "Did you include both srcjar and a sources jar in the jars attribute?: %s" % source_jars)
        return struct(class_jar = jars[0], source_jar = source_jars[0] if len(source_jars) == 1 else None, ijar = None)

def kt_jvm_import_impl(ctx):
    if bool(ctx.attr.jars) and bool(ctx.attr.jar):
        fail("Cannot use both jars= and jar= attribute.  Prefer jar=")

    artifact = _unify_jars(ctx)
    kt_info = _KtJvmInfo(
        module_name = _utils.derive_module_name(ctx),
        exported_compiler_plugins = depset(getattr(ctx.attr, "exported_compiler_plugins", [])),
        outputs = struct(
            jars = [artifact],
        ),
    )
    return struct(
        kt = kt_info,
        providers = [
            DefaultInfo(
                files = depset(direct = [artifact.class_jar]),
                runfiles = ctx.runfiles(files = [artifact.class_jar, artifact.source_jar]),
            ),
            JavaInfo(
                output_jar = artifact.class_jar,
                compile_jar = artifact.class_jar,
                source_jar = artifact.source_jar,
                runtime_deps = [dep[JavaInfo] for dep in ctx.attr.runtime_deps if JavaInfo in dep],
                deps = [dep[JavaInfo] for dep in ctx.attr.deps if JavaInfo in dep],
                exports = [d[JavaInfo] for d in getattr(ctx.attr, "exports", [])],
                neverlink = getattr(ctx.attr, "neverlink", False),
            ),
            kt_info,
        ],
    )

def kt_jvm_library_impl(ctx):
    if ctx.attr.neverlink and ctx.attr.runtime_deps:
        fail("runtime_deps and neverlink is nonsensical.", attr = "runtime_deps")

    if not ctx.attr.srcs and len(ctx.attr.deps) > 0:
        fail(
            "deps without srcs in invalid." +
            "\nTo add runtime dependencies use runtime_deps." +
            "\nTo export libraries use exports.",
            attr = "deps",
        )

    return _make_providers(
        ctx,
        _kt_jvm_produce_jar_actions(ctx, "kt_jvm_library") if ctx.attr.srcs else export_only_providers(
            ctx = ctx,
            actions = ctx.actions,
            outputs = ctx.outputs,
            attr = ctx.attr,
        ),
    )

def kt_jvm_binary_impl(ctx):
    providers = _kt_jvm_produce_jar_actions(ctx, "kt_jvm_binary")
    _write_launcher_action(
        ctx,
        providers.java.transitive_runtime_jars,
        ctx.attr.main_class,
        ctx.attr.jvm_flags,
    )
    if len(ctx.attr.srcs) == 0 and len(ctx.attr.deps) > 0:
        fail("deps without srcs in invalid. To add runtime classpath and resources, use runtime_deps.", attr = "deps")

    return _make_providers(
        ctx,
        providers,
        depset(
            order = "default",
            transitive = [providers.java.transitive_runtime_jars],
            direct = ctx.files._java_runtime,
        ),
    )

_SPLIT_STRINGS = [
    "src/test/java/",
    "src/test/kotlin/",
    "javatests/",
    "kotlin/",
    "java/",
    "test/",
]

def kt_jvm_junit_test_impl(ctx):
    providers = _kt_jvm_produce_jar_actions(ctx, "kt_jvm_test")
    runtime_jars = depset(ctx.files._bazel_test_runner, transitive = [providers.java.transitive_runtime_jars])

    test_class = ctx.attr.test_class

    # If no test_class, do a best-effort attempt to infer one.
    if not bool(ctx.attr.test_class):
        for file in ctx.files.srcs:
            if file.basename.split(".")[0] == ctx.attr.name:
                for splitter in _SPLIT_STRINGS:
                    elements = file.short_path.split(splitter, 1)
                    if len(elements) == 2:
                        test_class = elements[1].split(".")[0].replace("/", ".")

    _write_launcher_action(
        ctx,
        runtime_jars,
        main_class = ctx.attr.main_class,
        jvm_flags = [
            "-ea",
            "-Dbazel.test_suite=%s" % test_class,
        ] + ctx.attr.jvm_flags,
    )

    return _make_providers(
        ctx,
        providers,
        depset(
            order = "default",
            transitive = [runtime_jars],
            direct = ctx.files._java_runtime,
        ),
        # adds common test variables, including TEST_WORKSPACE.
        testing.TestEnvironment({}),
    )

_KtCompilerPluginInfoDeps = provider(
    fields = {
        "compiler_libs": "list javainfo of a compiler library",
    },
)

def kt_compiler_deps_aspect_impl(target, ctx):
    "Collects the dependencies of a target and forward it to allow deshading the plugin jars."
    return _KtCompilerPluginInfoDeps(
        compiler_libs = [
            t[JavaInfo]
            for d in ["deps", "runtime_deps", "exports"]
            for t in getattr(ctx.rule.attr, d, [])
            if JavaInfo in t
        ],
    )

def kt_compiler_plugin_impl(ctx):
    plugin_id = ctx.attr.id
    options = []
    for (k, v) in ctx.attr.options.items():
        if "=" in k:
            fail("kt_compiler_plugin options keys cannot contain the = symbol")
        options.append(struct(id = plugin_id, value = "%s=%s" % (k, v)))

    jars = []
    compiler_lib_infos = []
    for t in [t for t in ctx.attr.deps]:
        compiler_lib_infos.extend(t[_KtCompilerPluginInfoDeps].compiler_libs)
        ji = t[JavaInfo]
        if ji.outputs:  # can be None, according to the docs.
            for jar_output in ji.outputs.jars:
                jar = jar_output.class_jar
                if ctx.attr.target_embedded_compiler:
                    jars.append(jarjar_action(
                        actions = ctx.actions,
                        jarjar = ctx.executable._jarjar,
                        label = ctx.label,
                        rules = ctx.file._jetbrains_deshade_rules,
                        input = jar,
                        output = ctx.actions.declare_file("%s_deshaded_%s" % (ctx.label.name, jar.basename)),
                    ))
                else:
                    jars.append(jar)

    if not jars:
        fail("Unable to locate plugin jars. Ensure they are exported or the direct result of a library in the deps.")

    if not (ctx.attr.compile_phase or ctx.attr.stubs_phase):
        fail("Plugin must execute during in one or more phases: stubs_phase, compile_phase")

    info = java_common.merge([
        JavaInfo(output_jar = d, compile_jar = d, deps = compiler_lib_infos)
        for d in jars
    ])

    return [
        DefaultInfo(files = depset(jars)),
        java_common.merge([]),
        _KtCompilerPluginInfo(
            plugin_jars = jars,
            classpath = depset(jars, transitive = [i.transitive_deps for i in compiler_lib_infos]),
            options = options,
            stubs = ctx.attr.stubs_phase,
            compile = ctx.attr.compile_phase,
        ),
    ]
<commit_msg>Fix test class checking to break when it finds a match. (#464)<commit_after># Copyright 2018 The Bazel Authors. 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.
load(
    "//kotlin/internal/jvm:compile.bzl",
    "export_only_providers",
    _kt_jvm_produce_jar_actions = "kt_jvm_produce_jar_actions",
)
load(
    "//kotlin/internal:defs.bzl",
    _KtCompilerPluginInfo = "KtCompilerPluginInfo",
    _KtJvmInfo = "KtJvmInfo",
)
load(
    "//kotlin/internal/utils:utils.bzl",
    _utils = "utils",
)
load("//third_party:jarjar.bzl", "jarjar_action")

def _make_providers(ctx, providers, transitive_files = depset(order = "default"), *additional_providers):
    return struct(
        kt = providers.kt,
        providers = [
            providers.java,
            providers.kt,
            DefaultInfo(
                files = depset([ctx.outputs.jar, ctx.outputs.jdeps]),
                runfiles = ctx.runfiles(
                    # explicitly include data files, otherwise they appear to be missing
                    files = ctx.files.data,
                    transitive_files = transitive_files,
                    # continue to use collect_default until proper transitive data collecting is
                    # implmented.
                    collect_default = True,
                ),
            ),
        ] + list(additional_providers),
    )

def _write_launcher_action(ctx, rjars, main_class, jvm_flags, args = "", wrapper_preamble = ""):
    """Macro that writes out a launcher script shell script.
      Args:
        rjars: All of the runtime jars required to launch this java target.
        main_class: the main class to launch.
        jvm_flags: The flags that should be passed to the jvm.
        args: Args that should be passed to the Binary.
    """
    classpath = ":".join(["${RUNPATH}%s" % (j.short_path) for j in rjars.to_list()])
    jvm_flags = " ".join([ctx.expand_location(f, ctx.attr.data) for f in jvm_flags])
    template = ctx.attr._java_stub_template.files.to_list()[0]
    javabin = "JAVABIN=" + str(ctx.attr._java_runtime[java_common.JavaRuntimeInfo].java_executable_exec_path)

    ctx.actions.expand_template(
        template = template,
        output = ctx.outputs.executable,
        substitutions = {
            "%classpath%": classpath,
            "%java_start_class%": main_class,
            "%javabin%": javabin,
            "%jvm_flags%": jvm_flags,
            "%set_jacoco_metadata%": "",
            "%set_jacoco_main_class%": "",
            "%set_jacoco_java_runfiles_root%": "",
            "%set_java_coverage_new_implementation%": "",
            "%workspace_prefix%": ctx.workspace_name + "/",
        },
        is_executable = True,
    )

def _is_source_jar_stub(jar):
    """Workaround for intellij plugin expecting a source jar"""
    return jar.path.endswith("third_party/empty.jar")

def _unify_jars(ctx):
    if bool(ctx.attr.jar):
        return struct(class_jar = ctx.file.jar, source_jar = ctx.file.srcjar, ijar = None)
    else:
        # Legacy handling.
        jars = []
        source_jars = []
        if (ctx.file.srcjar and not ("%s" % ctx.file.srcjar.path).endswith("third_party/empty.jar")):
            source_jars += [ctx.file.srcjar]

        # There must be a single jar jar and it can either be a filegroup or a JavaInfo.
        for jar in ctx.attr.jars:
            # If a JavaInfo is available it's because it was picked up from a `maven_jar` style attribute -- e.g.,
            # @com_google_guava_guava//jar. so the transitive_compile_jars or the transitive_runtime_jars should not be
            # visited -- descending into these results in ijars entering the graph.
            if JavaInfo in jar:
                jars += jar[JavaInfo].full_compile_jars.to_list()
                source_jars += jar[JavaInfo].transitive_source_jars.to_list()
            else:
                # this branch occurs when the attr was a filegroup.
                for file in jar.files.to_list():
                    if file.basename.endswith("-sources.jar"):
                        source_jars.append(file)
                    elif file.basename.endswith(".jar"):
                        jars.append(file)
                    else:
                        fail("a jar pointing to a filegroup must either end with -sources.jar or .jar")

        if len(jars) > 1:
            fail("Got more than one jar, this is an error create an issue: %s" % jars)
        if len(source_jars) > 1:
            fail("Got more than one source jar. " +
                 "Did you include both srcjar and a sources jar in the jars attribute?: %s" % source_jars)
        return struct(class_jar = jars[0], source_jar = source_jars[0] if len(source_jars) == 1 else None, ijar = None)

def kt_jvm_import_impl(ctx):
    if bool(ctx.attr.jars) and bool(ctx.attr.jar):
        fail("Cannot use both jars= and jar= attribute.  Prefer jar=")

    artifact = _unify_jars(ctx)
    kt_info = _KtJvmInfo(
        module_name = _utils.derive_module_name(ctx),
        exported_compiler_plugins = depset(getattr(ctx.attr, "exported_compiler_plugins", [])),
        outputs = struct(
            jars = [artifact],
        ),
    )
    return struct(
        kt = kt_info,
        providers = [
            DefaultInfo(
                files = depset(direct = [artifact.class_jar]),
                runfiles = ctx.runfiles(files = [artifact.class_jar, artifact.source_jar]),
            ),
            JavaInfo(
                output_jar = artifact.class_jar,
                compile_jar = artifact.class_jar,
                source_jar = artifact.source_jar,
                runtime_deps = [dep[JavaInfo] for dep in ctx.attr.runtime_deps if JavaInfo in dep],
                deps = [dep[JavaInfo] for dep in ctx.attr.deps if JavaInfo in dep],
                exports = [d[JavaInfo] for d in getattr(ctx.attr, "exports", [])],
                neverlink = getattr(ctx.attr, "neverlink", False),
            ),
            kt_info,
        ],
    )

def kt_jvm_library_impl(ctx):
    if ctx.attr.neverlink and ctx.attr.runtime_deps:
        fail("runtime_deps and neverlink is nonsensical.", attr = "runtime_deps")

    if not ctx.attr.srcs and len(ctx.attr.deps) > 0:
        fail(
            "deps without srcs in invalid." +
            "\nTo add runtime dependencies use runtime_deps." +
            "\nTo export libraries use exports.",
            attr = "deps",
        )

    return _make_providers(
        ctx,
        _kt_jvm_produce_jar_actions(ctx, "kt_jvm_library") if ctx.attr.srcs else export_only_providers(
            ctx = ctx,
            actions = ctx.actions,
            outputs = ctx.outputs,
            attr = ctx.attr,
        ),
    )

def kt_jvm_binary_impl(ctx):
    providers = _kt_jvm_produce_jar_actions(ctx, "kt_jvm_binary")
    _write_launcher_action(
        ctx,
        providers.java.transitive_runtime_jars,
        ctx.attr.main_class,
        ctx.attr.jvm_flags,
    )
    if len(ctx.attr.srcs) == 0 and len(ctx.attr.deps) > 0:
        fail("deps without srcs in invalid. To add runtime classpath and resources, use runtime_deps.", attr = "deps")

    return _make_providers(
        ctx,
        providers,
        depset(
            order = "default",
            transitive = [providers.java.transitive_runtime_jars],
            direct = ctx.files._java_runtime,
        ),
    )

_SPLIT_STRINGS = [
    "src/test/java/",
    "src/test/kotlin/",
    "javatests/",
    "kotlin/",
    "java/",
    "test/",
]

def kt_jvm_junit_test_impl(ctx):
    providers = _kt_jvm_produce_jar_actions(ctx, "kt_jvm_test")
    runtime_jars = depset(ctx.files._bazel_test_runner, transitive = [providers.java.transitive_runtime_jars])

    test_class = ctx.attr.test_class

    # If no test_class, do a best-effort attempt to infer one.
    if not bool(ctx.attr.test_class):
        for file in ctx.files.srcs:
            if file.basename.split(".")[0] == ctx.attr.name:
                for splitter in _SPLIT_STRINGS:
                    elements = file.short_path.split(splitter, 1)
                    if len(elements) == 2:
                        test_class = elements[1].split(".")[0].replace("/", ".")
                        break

    _write_launcher_action(
        ctx,
        runtime_jars,
        main_class = ctx.attr.main_class,
        jvm_flags = [
            "-ea",
            "-Dbazel.test_suite=%s" % test_class,
        ] + ctx.attr.jvm_flags,
    )

    return _make_providers(
        ctx,
        providers,
        depset(
            order = "default",
            transitive = [runtime_jars],
            direct = ctx.files._java_runtime,
        ),
        # adds common test variables, including TEST_WORKSPACE.
        testing.TestEnvironment({}),
    )

_KtCompilerPluginInfoDeps = provider(
    fields = {
        "compiler_libs": "list javainfo of a compiler library",
    },
)

def kt_compiler_deps_aspect_impl(target, ctx):
    "Collects the dependencies of a target and forward it to allow deshading the plugin jars."
    return _KtCompilerPluginInfoDeps(
        compiler_libs = [
            t[JavaInfo]
            for d in ["deps", "runtime_deps", "exports"]
            for t in getattr(ctx.rule.attr, d, [])
            if JavaInfo in t
        ],
    )

def kt_compiler_plugin_impl(ctx):
    plugin_id = ctx.attr.id
    options = []
    for (k, v) in ctx.attr.options.items():
        if "=" in k:
            fail("kt_compiler_plugin options keys cannot contain the = symbol")
        options.append(struct(id = plugin_id, value = "%s=%s" % (k, v)))

    jars = []
    compiler_lib_infos = []
    for t in [t for t in ctx.attr.deps]:
        compiler_lib_infos.extend(t[_KtCompilerPluginInfoDeps].compiler_libs)
        ji = t[JavaInfo]
        if ji.outputs:  # can be None, according to the docs.
            for jar_output in ji.outputs.jars:
                jar = jar_output.class_jar
                if ctx.attr.target_embedded_compiler:
                    jars.append(jarjar_action(
                        actions = ctx.actions,
                        jarjar = ctx.executable._jarjar,
                        label = ctx.label,
                        rules = ctx.file._jetbrains_deshade_rules,
                        input = jar,
                        output = ctx.actions.declare_file("%s_deshaded_%s" % (ctx.label.name, jar.basename)),
                    ))
                else:
                    jars.append(jar)

    if not jars:
        fail("Unable to locate plugin jars. Ensure they are exported or the direct result of a library in the deps.")

    if not (ctx.attr.compile_phase or ctx.attr.stubs_phase):
        fail("Plugin must execute during in one or more phases: stubs_phase, compile_phase")

    info = java_common.merge([
        JavaInfo(output_jar = d, compile_jar = d, deps = compiler_lib_infos)
        for d in jars
    ])

    return [
        DefaultInfo(files = depset(jars)),
        java_common.merge([]),
        _KtCompilerPluginInfo(
            plugin_jars = jars,
            classpath = depset(jars, transitive = [i.transitive_deps for i in compiler_lib_infos]),
            options = options,
            stubs = ctx.attr.stubs_phase,
            compile = ctx.attr.compile_phase,
        ),
    ]
<|endoftext|>


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

from typing import Callable, Union, Tuple, Dict
from textwrap import indent
from json import dumps, encoder
from logging import DEBUG

from sklearn.tree.tree import DecisionTreeClassifier \
    as DecisionTreeClassifierClass

from sklearn_porter.estimator.EstimatorApiABC import EstimatorApiABC
from sklearn_porter.estimator.EstimatorBase import EstimatorBase
from sklearn_porter.enums import Method, Language, Template
from sklearn_porter.utils import get_logger

L = get_logger(__name__)

class DecisionTreeClassifier(EstimatorBase, EstimatorApiABC):
    """
    Extract model data and port a DecisionTreeClassifier classifier.

    See also
    --------
    http://scikit-learn.org/stable/modules/generated/sklearn.tree.DecisionTreeClassifier.html
    """
    FULL_SUPPORT = {  # shorthandler for the support matrix
        Method.PREDICT: {
            Template.COMBINED,
            Template.ATTACHED,
            Template.EXPORTED
        }
    }
    support = {
        Language.C: FULL_SUPPORT,
        Language.GO: FULL_SUPPORT,
        Language.JAVA: FULL_SUPPORT,
        Language.JS: FULL_SUPPORT,
        Language.PHP: FULL_SUPPORT,
        Language.RUBY: FULL_SUPPORT
    }

    estimator = None  # type: DecisionTreeClassifierClass

    def __init__(self, estimator: DecisionTreeClassifierClass):
        super().__init__(estimator)
        L.info('Create specific estimator `%s`.', self.estimator_name)
        est = self.estimator  # alias

        # Extract and save meta information:
        self.meta_info = dict(
            n_features=est.n_features_,
            n_classes=len(est.tree_.value.tolist()[0][0]),
        )
        L.info('Extracted meta information (keys only): {}'.format(
            self.meta_info.keys()))
        if L.isEnabledFor(DEBUG):
            dump = EstimatorBase._dump_dict(self.meta_info)
            L.debug('Extracted meta information:\n{}'.format(dump))

        # Extract and save model data:
        self.model_data = dict(
            lefts=est.tree_.children_left.tolist(),
            rights=est.tree_.children_right.tolist(),
            thresholds=est.tree_.threshold.tolist(),
            indices=est.tree_.feature.tolist(),
            classes=[[int(c) for c in l[0]] for l in est.tree_.value.tolist()],
        )
        L.info('Extracted model data (keys only): {}'.format(
            self.model_data.keys()))
        if L.isEnabledFor(DEBUG):
            dump = EstimatorBase._dump_dict(self.model_data)
            L.debug('Extracted model data:\n{}'.format(dump))

    def port(
            self,
            method: Method,
            language: Language,
            template: Template,
            **kwargs
    ) -> Union[str, Tuple[str, str]]:
        """
        Port an estimator.

        Parameters
        ----------
        method : Method
            The required method.
        language : Language
            The required language.
        template : Template
            The required template.
        kwargs

        Returns
        -------
        The ported estimator.
        """

        super().check_arguments(method, language, template)

        converter = kwargs.get('converter')

        # Placeholders:
        placeholders = dict(
            class_name=kwargs.get('class_name'),
            method_name=kwargs.get('method_name'),
        )
        placeholders.update(self.meta_info)

        # Load templates:
        tpls = self._load_templates(language.value.KEY)

        if template == Template.EXPORTED:
            output = str(tpls.get('exported.class').format(**placeholders))
            converter = kwargs.get('converter')
            encoder.FLOAT_REPR = lambda o: converter(o)
            model_data = dumps(self.model_data, sort_keys=True)
            return output, model_data

        # Pick templates:
        tpl_int = tpls.get('int')
        tpl_double = tpls.get('double')
        tpl_arr_1 = tpls.get('arr[]')
        tpl_arr_2 = tpls.get('arr[][]')
        tpl_in_brackets = tpls.get('in_brackets')

        # Make contents:
        lefts_val = list(map(str, self.model_data['lefts']))
        lefts_str = tpl_arr_1.format(type=tpl_int, name='lefts',
                                     values=', '.join(lefts_val),
                                     n=len(lefts_val))

        rights_val = list(map(str, self.model_data['rights']))
        rights_str = tpl_arr_1.format(type=tpl_int, name='rights',
                                      values=', '.join(rights_val),
                                      n=len(rights_val))

        thresholds_val = list(map(converter, self.model_data['thresholds']))
        thresholds_str = tpl_arr_1.format(type=tpl_double, name='thresholds',
                                          values=', '.join(thresholds_val),
                                          n=len(thresholds_val))

        indices_val = list(map(str, self.model_data['indices']))
        indices_str = tpl_arr_1.format(type=tpl_int, name='indices',
                                       values=', '.join(indices_val),
                                       n=len(indices_val))

        classes_val = [list(map(str, e)) for e in self.model_data['classes']]
        n, m = len(classes_val), self.meta_info.get('n_classes')
        classes_str = [', '.join(e) for e in classes_val]
        classes_str = ', '.join([tpl_in_brackets.format(e)
                                 for e in classes_str])
        classes_str = tpl_arr_2.format(type=tpl_int, name='classes',
                                       values=classes_str, n=n, m=m)

        placeholders.update(dict(
            lefts=lefts_str,
            rights=rights_str,
            thresholds=thresholds_str,
            indices=indices_str,
            classes=classes_str,
        ))

        if template == Template.ATTACHED:
            return tpls.get('attached.class').format(**placeholders)

        if template == Template.COMBINED:

            # Pick templates:
            tpl_indent = tpls.get('indent')
            tpl_method = tpls.get('combined.method')
            tpl_class = tpls.get('combined.class')

            # Make tree:
            out_tree = self._create_tree(tpls, language.value.KEY, converter)
            out_tree = indent(out_tree, 1 * tpl_indent)

            # Make method:
            placeholders.update(dict(tree=out_tree))
            n_indents = 1 if language in [
                Language.JAVA,
                Language.JS,
                Language.PHP,
                Language.RUBY
            ] else 0
            tpl_method = indent(tpl_method, n_indents * tpl_indent)
            tpl_method = tpl_method[(n_indents * len(tpl_indent)):]
            out_method = tpl_method.format(**placeholders)

            # Make class:
            placeholders.update(dict(method=out_method))
            out_class = tpl_class.format(**placeholders)

            return out_class

    def _create_tree(
            self,
            tpls: Dict[str, str],
            language: str,
            converter: Callable[[object], str]
    ):
        """
        Build a decision tree.

        Parameters
        ----------
        tpls : Dict[str, str]
            All relevant templates.
        language : str
            The required language.
        converter : Callable[[object], str]
            The number converter.

        Returns
        -------
        A tree of a DecisionTreeClassifier.
        """
        feature_indices = []
        for i in self.model_data['indices']:
            n_features = self.meta_info['n_features']
            if n_features > 1 or (n_features == 1 and i >= 0):
                feature_indices.append([str(j) for j in range(n_features)][i])

        n_indents = 1 if language in {'java', 'js', 'php', 'ruby'} else 0
        return self._create_branch(
            tpls, language, converter,
            self.model_data['lefts'],
            self.model_data['rights'],
            self.model_data['thresholds'],
            self.model_data['classes'],
            feature_indices, 0, n_indents)

    def _create_branch(
            self,
            tpls: dict,
            language: str,
            converter: Callable[[object], str],
            left_nodes: list,
            right_nodes: list,
            threshold: list,
            value: list,
            features: list,
            node: int,
            depth: int
    ):
        """
        The ported single tree as function or method.

        Parameters
        ----------
        tpls : Dict[str, str]
            All relevant templates.
        language
            The required language.
        converter
            The number converter.
        left_nodes : list
            The left children node.
        right_nodes : list
            The left children node.
        threshold : list
            The decision thresholds.
        value : list
            The label or class.
        features : list
            The feature values.
        node : list
            The current node.
        depth : list
            The tree depth.

        Returns
        -------
        A single branch of a DecisionTreeClassifier.
        """
        out = ''
        temp_indent = tpls.get('indent')
        if threshold[node] != -2.:

            out += '\n'
            tpl = tpls.get('if')
            tpl = indent(tpl, depth * temp_indent)
            val_1 = 'features[{}]'.format(features[node])
            if language == 'php':
                val_1 = '$' + val_1
            val_2 = converter(threshold[node])
            out += tpl.format(val_1, '<=', val_2)

            if left_nodes[node] != -1.:
                out += self._create_branch(
                    tpls, language, converter, left_nodes, right_nodes,
                    threshold, value, features, left_nodes[node], depth + 1)

            out += '\n'
            tpl = tpls.get('else')
            tpl = indent(tpl, depth * temp_indent)
            out += tpl

            if right_nodes[node] != -1.:
                out += self._create_branch(
                    tpls, language, converter, left_nodes, right_nodes,
                    threshold, value, features, right_nodes[node], depth + 1)

            out += '\n'
            tpl = tpls.get('endif')
            tpl = indent(tpl, depth * temp_indent)
            out += tpl
        else:
            clazzes = []
            tpl = 'classes[{0}] = {1}'
            if language == 'php':
                tpl = '$' + tpl
            tpl = indent(tpl, depth * temp_indent)

            for i, rate in enumerate(value[node]):
                clazz = tpl.format(i, rate)
                clazz = '\n' + clazz
                clazzes.append(clazz)

            tpl = tpls.get('join')
            out += tpl.join(clazzes) + tpl
        return out
<commit_msg>feature/oop-api-refactoring: Get the most compact JSON representation with `(',', ':')` for the model data dump<commit_after># -*- coding: utf-8 -*-

from typing import Callable, Union, Tuple, Dict
from textwrap import indent
from json import dumps, encoder
from logging import DEBUG

from sklearn.tree.tree import DecisionTreeClassifier \
    as DecisionTreeClassifierClass

from sklearn_porter.estimator.EstimatorApiABC import EstimatorApiABC
from sklearn_porter.estimator.EstimatorBase import EstimatorBase
from sklearn_porter.enums import Method, Language, Template
from sklearn_porter.utils import get_logger

L = get_logger(__name__)

class DecisionTreeClassifier(EstimatorBase, EstimatorApiABC):
    """
    Extract model data and port a DecisionTreeClassifier classifier.

    See also
    --------
    http://scikit-learn.org/stable/modules/generated/sklearn.tree.DecisionTreeClassifier.html
    """
    FULL_SUPPORT = {  # shorthandler for the support matrix
        Method.PREDICT: {
            Template.COMBINED,
            Template.ATTACHED,
            Template.EXPORTED
        }
    }
    support = {
        Language.C: FULL_SUPPORT,
        Language.GO: FULL_SUPPORT,
        Language.JAVA: FULL_SUPPORT,
        Language.JS: FULL_SUPPORT,
        Language.PHP: FULL_SUPPORT,
        Language.RUBY: FULL_SUPPORT
    }

    estimator = None  # type: DecisionTreeClassifierClass

    def __init__(self, estimator: DecisionTreeClassifierClass):
        super().__init__(estimator)
        L.info('Create specific estimator `%s`.', self.estimator_name)
        est = self.estimator  # alias

        # Extract and save meta information:
        self.meta_info = dict(
            n_features=est.n_features_,
            n_classes=len(est.tree_.value.tolist()[0][0]),
        )
        L.info('Extracted meta information (keys only): {}'.format(
            self.meta_info.keys()))
        if L.isEnabledFor(DEBUG):
            dump = EstimatorBase._dump_dict(self.meta_info)
            L.debug('Extracted meta information:\n{}'.format(dump))

        # Extract and save model data:
        self.model_data = dict(
            lefts=est.tree_.children_left.tolist(),
            rights=est.tree_.children_right.tolist(),
            thresholds=est.tree_.threshold.tolist(),
            indices=est.tree_.feature.tolist(),
            classes=[[int(c) for c in l[0]] for l in est.tree_.value.tolist()],
        )
        L.info('Extracted model data (keys only): {}'.format(
            self.model_data.keys()))
        if L.isEnabledFor(DEBUG):
            dump = EstimatorBase._dump_dict(self.model_data)
            L.debug('Extracted model data:\n{}'.format(dump))

    def port(
            self,
            method: Method,
            language: Language,
            template: Template,
            **kwargs
    ) -> Union[str, Tuple[str, str]]:
        """
        Port an estimator.

        Parameters
        ----------
        method : Method
            The required method.
        language : Language
            The required language.
        template : Template
            The required template.
        kwargs

        Returns
        -------
        The ported estimator.
        """

        super().check_arguments(method, language, template)

        converter = kwargs.get('converter')

        # Placeholders:
        placeholders = dict(
            class_name=kwargs.get('class_name'),
            method_name=kwargs.get('method_name'),
        )
        placeholders.update(self.meta_info)

        # Load templates:
        tpls = self._load_templates(language.value.KEY)

        if template == Template.EXPORTED:
            output = str(tpls.get('exported.class').format(**placeholders))
            converter = kwargs.get('converter')
            encoder.FLOAT_REPR = lambda o: converter(o)
            model_data = dumps(self.model_data, separators=(',', ':'))
            return output, model_data

        # Pick templates:
        tpl_int = tpls.get('int')
        tpl_double = tpls.get('double')
        tpl_arr_1 = tpls.get('arr[]')
        tpl_arr_2 = tpls.get('arr[][]')
        tpl_in_brackets = tpls.get('in_brackets')

        # Make contents:
        lefts_val = list(map(str, self.model_data['lefts']))
        lefts_str = tpl_arr_1.format(type=tpl_int, name='lefts',
                                     values=', '.join(lefts_val),
                                     n=len(lefts_val))

        rights_val = list(map(str, self.model_data['rights']))
        rights_str = tpl_arr_1.format(type=tpl_int, name='rights',
                                      values=', '.join(rights_val),
                                      n=len(rights_val))

        thresholds_val = list(map(converter, self.model_data['thresholds']))
        thresholds_str = tpl_arr_1.format(type=tpl_double, name='thresholds',
                                          values=', '.join(thresholds_val),
                                          n=len(thresholds_val))

        indices_val = list(map(str, self.model_data['indices']))
        indices_str = tpl_arr_1.format(type=tpl_int, name='indices',
                                       values=', '.join(indices_val),
                                       n=len(indices_val))

        classes_val = [list(map(str, e)) for e in self.model_data['classes']]
        n, m = len(classes_val), self.meta_info.get('n_classes')
        classes_str = [', '.join(e) for e in classes_val]
        classes_str = ', '.join([tpl_in_brackets.format(e)
                                 for e in classes_str])
        classes_str = tpl_arr_2.format(type=tpl_int, name='classes',
                                       values=classes_str, n=n, m=m)

        placeholders.update(dict(
            lefts=lefts_str,
            rights=rights_str,
            thresholds=thresholds_str,
            indices=indices_str,
            classes=classes_str,
        ))

        if template == Template.ATTACHED:
            return tpls.get('attached.class').format(**placeholders)

        if template == Template.COMBINED:

            # Pick templates:
            tpl_indent = tpls.get('indent')
            tpl_method = tpls.get('combined.method')
            tpl_class = tpls.get('combined.class')

            # Make tree:
            out_tree = self._create_tree(tpls, language.value.KEY, converter)
            out_tree = indent(out_tree, 1 * tpl_indent)

            # Make method:
            placeholders.update(dict(tree=out_tree))
            n_indents = 1 if language in [
                Language.JAVA,
                Language.JS,
                Language.PHP,
                Language.RUBY
            ] else 0
            tpl_method = indent(tpl_method, n_indents * tpl_indent)
            tpl_method = tpl_method[(n_indents * len(tpl_indent)):]
            out_method = tpl_method.format(**placeholders)

            # Make class:
            placeholders.update(dict(method=out_method))
            out_class = tpl_class.format(**placeholders)

            return out_class

    def _create_tree(
            self,
            tpls: Dict[str, str],
            language: str,
            converter: Callable[[object], str]
    ):
        """
        Build a decision tree.

        Parameters
        ----------
        tpls : Dict[str, str]
            All relevant templates.
        language : str
            The required language.
        converter : Callable[[object], str]
            The number converter.

        Returns
        -------
        A tree of a DecisionTreeClassifier.
        """
        feature_indices = []
        for i in self.model_data['indices']:
            n_features = self.meta_info['n_features']
            if n_features > 1 or (n_features == 1 and i >= 0):
                feature_indices.append([str(j) for j in range(n_features)][i])

        n_indents = 1 if language in {'java', 'js', 'php', 'ruby'} else 0
        return self._create_branch(
            tpls, language, converter,
            self.model_data['lefts'],
            self.model_data['rights'],
            self.model_data['thresholds'],
            self.model_data['classes'],
            feature_indices, 0, n_indents)

    def _create_branch(
            self,
            tpls: dict,
            language: str,
            converter: Callable[[object], str],
            left_nodes: list,
            right_nodes: list,
            threshold: list,
            value: list,
            features: list,
            node: int,
            depth: int
    ):
        """
        The ported single tree as function or method.

        Parameters
        ----------
        tpls : Dict[str, str]
            All relevant templates.
        language
            The required language.
        converter
            The number converter.
        left_nodes : list
            The left children node.
        right_nodes : list
            The left children node.
        threshold : list
            The decision thresholds.
        value : list
            The label or class.
        features : list
            The feature values.
        node : list
            The current node.
        depth : list
            The tree depth.

        Returns
        -------
        A single branch of a DecisionTreeClassifier.
        """
        out = ''
        temp_indent = tpls.get('indent')
        if threshold[node] != -2.:

            out += '\n'
            tpl = tpls.get('if')
            tpl = indent(tpl, depth * temp_indent)
            val_1 = 'features[{}]'.format(features[node])
            if language == 'php':
                val_1 = '$' + val_1
            val_2 = converter(threshold[node])
            out += tpl.format(val_1, '<=', val_2)

            if left_nodes[node] != -1.:
                out += self._create_branch(
                    tpls, language, converter, left_nodes, right_nodes,
                    threshold, value, features, left_nodes[node], depth + 1)

            out += '\n'
            tpl = tpls.get('else')
            tpl = indent(tpl, depth * temp_indent)
            out += tpl

            if right_nodes[node] != -1.:
                out += self._create_branch(
                    tpls, language, converter, left_nodes, right_nodes,
                    threshold, value, features, right_nodes[node], depth + 1)

            out += '\n'
            tpl = tpls.get('endif')
            tpl = indent(tpl, depth * temp_indent)
            out += tpl
        else:
            clazzes = []
            tpl = 'classes[{0}] = {1}'
            if language == 'php':
                tpl = '$' + tpl
            tpl = indent(tpl, depth * temp_indent)

            for i, rate in enumerate(value[node]):
                clazz = tpl.format(i, rate)
                clazz = '\n' + clazz
                clazzes.append(clazz)

            tpl = tpls.get('join')
            out += tpl.join(clazzes) + tpl
        return out
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright 2017 Ricequant, Inc
#
# 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.

import datetime

import numpy as np
import numpy.lib.recfunctions as nprf

from rqalpha.interface import AbstractEventSource
from rqalpha.events import Event, EVENT
from rqalpha.environment import Environment
from rqalpha.utils import get_account_type
from rqalpha.utils.exception import CustomException, CustomError, patch_user_exc
from rqalpha.utils.datetime_func import convert_int_to_datetime, convert_date_to_date_int, convert_dt_to_int
from rqalpha.const import ACCOUNT_TYPE
from rqalpha.utils.i18n import gettext as _
from rqalpha.model.tick import Tick

ONE_MINUTE = datetime.timedelta(minutes=1)

class SimulationEventSource(AbstractEventSource):
    def __init__(self, env, account_list):
        self._env = env
        self._account_list = account_list
        self._universe_changed = False
        Environment.get_instance().event_bus.add_listener(EVENT.POST_UNIVERSE_CHANGED, self._on_universe_changed)

    def _on_universe_changed(self, event):
        self._universe_changed = True

    def _get_universe(self):
        universe = Environment.get_instance().get_universe()
        if len(universe) == 0 and ACCOUNT_TYPE.STOCK not in self._account_list:
            error = CustomError()
            error.set_msg("Current universe is empty. Please use subscribe function before trade")
            raise patch_user_exc(CustomException(error))
        return universe

    # [BEGIN] minute event helper
    @staticmethod
    def _get_stock_trading_minutes(trading_date):
        trading_minutes = set()
        current_dt = datetime.datetime.combine(trading_date, datetime.time(9, 31))
        am_end_dt = current_dt.replace(hour=11, minute=30)
        pm_start_dt = current_dt.replace(hour=13, minute=1)
        pm_end_dt = current_dt.replace(hour=15, minute=0)
        delta_minute = datetime.timedelta(minutes=1)
        while current_dt <= am_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute

        current_dt = pm_start_dt
        while current_dt <= pm_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute
        return trading_minutes

    def _get_future_trading_minutes(self, trading_date):
        trading_minutes = set()
        universe = self._get_universe()
        for order_book_id in universe:
            if get_account_type(order_book_id) == ACCOUNT_TYPE.STOCK:
                continue
            trading_minutes.update(self._env.data_proxy.get_trading_minutes_for(order_book_id, trading_date))
        return set([convert_int_to_datetime(minute) for minute in trading_minutes])

    def _get_trading_minutes(self, trading_date):
        trading_minutes = set()
        for account_type in self._account_list:
            if account_type == ACCOUNT_TYPE.STOCK:
                trading_minutes = trading_minutes.union(self._get_stock_trading_minutes(trading_date))
            elif account_type == ACCOUNT_TYPE.FUTURE:
                trading_minutes = trading_minutes.union(self._get_future_trading_minutes(trading_date))
        return sorted(list(trading_minutes))
    # [END] minute event helper

    def _get_merge_ticks(self, date, last_dt):
        self.today_ticks_cache = {}
        data_proxy = self._env.data_proxy
        self.order_book_id_map = {}

        # get today ticks
        for idx, order_book_id in enumerate(self._get_universe()):
            self.order_book_id_map[idx] = order_book_id
            ticks = data_proxy.get_ticks(order_book_id, date)
            if ticks is None:
                continue
            # append idx columns
            ticks = nprf.append_fields(ticks, 'order_book_id', np.full(len(ticks), idx), usemask=False)
            self.today_ticks_cache[order_book_id] = ticks

        # merge ticks arrays into one
        ticks_list = list(self.today_ticks_cache.values())
        if len(ticks_list) == 0:
            return []
        merge_ticks = np.concatenate(ticks_list)

        merge_ticks.sort(kind='mergesort', order=("date", "time"))

        time_arr = merge_ticks["date"].astype(np.uint64) * 1000000000 + merge_ticks["time"]
        last_idx = time_arr.searchsorted(convert_dt_to_int(last_dt) * 1000)
        merge_ticks = merge_ticks[last_idx:]

        return merge_ticks

    def events(self, start_date, end_date, frequency):
        if frequency == "1d":
            # 根据起始日期和结束日期，获取所有的交易日，然后再循环获取每一个交易日
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                date = day.to_pydatetime()
                dt_before_trading = date.replace(hour=0, minute=0)
                dt_bar = date.replace(hour=15, minute=0)
                dt_after_trading = date.replace(hour=15, minute=30)
                dt_settlement = date.replace(hour=17, minute=0)
                yield Event(EVENT.BEFORE_TRADING, calendar_dt=dt_before_trading, trading_dt=dt_before_trading)
                yield Event(EVENT.BAR, calendar_dt=dt_bar, trading_dt=dt_bar)

                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt_after_trading, trading_dt=dt_after_trading)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt_settlement, trading_dt=dt_settlement)
        elif frequency == '1m':
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                before_trading_flag = True
                date = day.to_pydatetime()
                last_dt = None
                done = False

                dt_before_day_trading = date.replace(hour=8, minute=30)

                while True:
                    if done:
                        break
                    exit_loop = True
                    trading_minutes = self._get_trading_minutes(date)
                    for calendar_dt in trading_minutes:
                        if last_dt is not None and calendar_dt < last_dt:
                            continue

                        if calendar_dt < dt_before_day_trading:
                            trading_dt = calendar_dt.replace(year=date.year,
                                                             month=date.month,
                                                             day=date.day)
                        else:
                            trading_dt = calendar_dt
                        if before_trading_flag:
                            before_trading_flag = False
                            before_trading_dt = trading_dt - datetime.timedelta(minutes=30)
                            yield Event(EVENT.BEFORE_TRADING, calendar_dt=before_trading_dt,
                                        trading_dt=before_trading_dt)
                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = calendar_dt
                            exit_loop = False
                            break
                        # yield handle bar
                        yield Event(EVENT.BAR, calendar_dt=calendar_dt, trading_dt=trading_dt)
                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        elif frequency == "tick":
            data_proxy = self._env.data_proxy
            for day in data_proxy.get_trading_dates(start_date, end_date):
                before_trading_flag = True
                date = day.to_pydatetime()
                last_dt = date
                done = False

                dt_before_day_trading = date.replace(hour=8, minute=30)

                while True:
                    if done:
                        break
                    exit_loop = True

                    merge_ticks = self._get_merge_ticks(date, last_dt)

                    for idx in range(len(merge_ticks)):
                        raw_tick = merge_ticks[idx]
                        order_book_id = self.order_book_id_map[raw_tick[-1]]

                        tick = Tick(order_book_id, raw_tick)
                        dt = tick.datetime

                        yield Event(EVENT.TICK, calendar_dt=dt, trading_dt=dt, tick=tick)

                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = dt
                            exit_loop = False
                            break

                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        else:
            raise NotImplementedError(_("Frequency {} is not support.").format(frequency))
<commit_msg>fix before_trading<commit_after># -*- coding: utf-8 -*-
#
# Copyright 2017 Ricequant, Inc
#
# 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.

import datetime

import numpy as np
import numpy.lib.recfunctions as nprf

from rqalpha.interface import AbstractEventSource
from rqalpha.events import Event, EVENT
from rqalpha.environment import Environment
from rqalpha.utils import get_account_type
from rqalpha.utils.exception import CustomException, CustomError, patch_user_exc
from rqalpha.utils.datetime_func import convert_int_to_datetime, convert_date_to_date_int, convert_dt_to_int
from rqalpha.const import ACCOUNT_TYPE
from rqalpha.utils.i18n import gettext as _
from rqalpha.model.tick import Tick

ONE_MINUTE = datetime.timedelta(minutes=1)

class SimulationEventSource(AbstractEventSource):
    def __init__(self, env, account_list):
        self._env = env
        self._account_list = account_list
        self._universe_changed = False
        Environment.get_instance().event_bus.add_listener(EVENT.POST_UNIVERSE_CHANGED, self._on_universe_changed)

    def _on_universe_changed(self, event):
        self._universe_changed = True

    def _get_universe(self):
        universe = Environment.get_instance().get_universe()
        if len(universe) == 0 and ACCOUNT_TYPE.STOCK not in self._account_list:
            error = CustomError()
            error.set_msg("Current universe is empty. Please use subscribe function before trade")
            raise patch_user_exc(CustomException(error))
        return universe

    # [BEGIN] minute event helper
    @staticmethod
    def _get_stock_trading_minutes(trading_date):
        trading_minutes = set()
        current_dt = datetime.datetime.combine(trading_date, datetime.time(9, 31))
        am_end_dt = current_dt.replace(hour=11, minute=30)
        pm_start_dt = current_dt.replace(hour=13, minute=1)
        pm_end_dt = current_dt.replace(hour=15, minute=0)
        delta_minute = datetime.timedelta(minutes=1)
        while current_dt <= am_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute

        current_dt = pm_start_dt
        while current_dt <= pm_end_dt:
            trading_minutes.add(current_dt)
            current_dt += delta_minute
        return trading_minutes

    def _get_future_trading_minutes(self, trading_date):
        trading_minutes = set()
        universe = self._get_universe()
        for order_book_id in universe:
            if get_account_type(order_book_id) == ACCOUNT_TYPE.STOCK:
                continue
            trading_minutes.update(self._env.data_proxy.get_trading_minutes_for(order_book_id, trading_date))
        return set([convert_int_to_datetime(minute) for minute in trading_minutes])

    def _get_trading_minutes(self, trading_date):
        trading_minutes = set()
        for account_type in self._account_list:
            if account_type == ACCOUNT_TYPE.STOCK:
                trading_minutes = trading_minutes.union(self._get_stock_trading_minutes(trading_date))
            elif account_type == ACCOUNT_TYPE.FUTURE:
                trading_minutes = trading_minutes.union(self._get_future_trading_minutes(trading_date))
        return sorted(list(trading_minutes))
    # [END] minute event helper

    def _get_merge_ticks(self, date, last_dt):
        self.today_ticks_cache = {}
        data_proxy = self._env.data_proxy
        self.order_book_id_map = {}

        # get today ticks
        for idx, order_book_id in enumerate(self._get_universe()):
            self.order_book_id_map[idx] = order_book_id
            ticks = data_proxy.get_ticks(order_book_id, date)
            if ticks is None:
                continue
            # append idx columns
            ticks = nprf.append_fields(ticks, 'order_book_id', np.full(len(ticks), idx), usemask=False)
            self.today_ticks_cache[order_book_id] = ticks

        # merge ticks arrays into one
        ticks_list = list(self.today_ticks_cache.values())
        if len(ticks_list) == 0:
            return []
        merge_ticks = np.concatenate(ticks_list)

        merge_ticks.sort(kind='mergesort', order=("date", "time"))

        time_arr = merge_ticks["date"].astype(np.uint64) * 1000000000 + merge_ticks["time"]
        last_idx = time_arr.searchsorted(convert_dt_to_int(last_dt) * 1000)
        merge_ticks = merge_ticks[last_idx:]

        return merge_ticks

    def events(self, start_date, end_date, frequency):
        if frequency == "1d":
            # 根据起始日期和结束日期，获取所有的交易日，然后再循环获取每一个交易日
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                date = day.to_pydatetime()
                dt_before_trading = date.replace(hour=0, minute=0)
                dt_bar = date.replace(hour=15, minute=0)
                dt_after_trading = date.replace(hour=15, minute=30)
                dt_settlement = date.replace(hour=17, minute=0)
                yield Event(EVENT.BEFORE_TRADING, calendar_dt=dt_before_trading, trading_dt=dt_before_trading)
                yield Event(EVENT.BAR, calendar_dt=dt_bar, trading_dt=dt_bar)

                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt_after_trading, trading_dt=dt_after_trading)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt_settlement, trading_dt=dt_settlement)
        elif frequency == '1m':
            for day in self._env.data_proxy.get_trading_dates(start_date, end_date):
                before_trading_flag = True
                date = day.to_pydatetime()
                last_dt = None
                done = False

                dt_before_day_trading = date.replace(hour=8, minute=30)

                while True:
                    if done:
                        break
                    exit_loop = True
                    trading_minutes = self._get_trading_minutes(date)
                    for calendar_dt in trading_minutes:
                        if last_dt is not None and calendar_dt < last_dt:
                            continue

                        if calendar_dt < dt_before_day_trading:
                            trading_dt = calendar_dt.replace(year=date.year,
                                                             month=date.month,
                                                             day=date.day)
                        else:
                            trading_dt = calendar_dt
                        if before_trading_flag:
                            before_trading_flag = False
                            before_trading_dt = trading_dt - datetime.timedelta(minutes=30)
                            yield Event(EVENT.BEFORE_TRADING, calendar_dt=before_trading_dt,
                                        trading_dt=before_trading_dt)
                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = calendar_dt
                            exit_loop = False
                            break
                        # yield handle bar
                        yield Event(EVENT.BAR, calendar_dt=calendar_dt, trading_dt=trading_dt)
                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        elif frequency == "tick":
            data_proxy = self._env.data_proxy
            for day in data_proxy.get_trading_dates(start_date, end_date):
                date = day.to_pydatetime()
                last_dt = date
                done = False

                before_trading_dt = date.replace(hour=8, minute=30)
                yield Event(EVENT.BEFORE_TRADING, calendar_dt=before_trading_dt,
                            trading_dt=before_trading_dt)

                while True:
                    if done:
                        break
                    exit_loop = True

                    merge_ticks = self._get_merge_ticks(date, last_dt)

                    for idx in range(len(merge_ticks)):
                        raw_tick = merge_ticks[idx]
                        order_book_id = self.order_book_id_map[raw_tick[-1]]

                        tick = Tick(order_book_id, raw_tick)
                        dt = tick.datetime

                        yield Event(EVENT.TICK, calendar_dt=dt, trading_dt=dt, tick=tick)

                        if self._universe_changed:
                            self._universe_changed = False
                            last_dt = dt
                            exit_loop = False
                            break

                    if exit_loop:
                        done = True

                dt = date.replace(hour=15, minute=30)
                yield Event(EVENT.AFTER_TRADING, calendar_dt=dt, trading_dt=dt)

                dt = date.replace(hour=17, minute=0)
                yield Event(EVENT.SETTLEMENT, calendar_dt=dt, trading_dt=dt)
        else:
            raise NotImplementedError(_("Frequency {} is not support.").format(frequency))
<|endoftext|>