<commit_before>#    Copyright 2015-2015 ARM Limited
#
# 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.
#

"""A library for asserting scheduler scenarios based on the
statistics aggregation framework"""

import re
import inspect
import trappy
from trappy.stats import SchedConf as sconf
from trappy.plotter.Utils import listify
from bart.sched.SchedAssert import SchedAssert
from bart.common import Utils

class SchedMultiAssert(object):

    """The primary focus of this class is to assert and verify
    predefined scheduler scenarios. This does not compare parameters
    across runs"""

    def __init__(self, run, topology, execnames):
        """Args:
                run (trappy.Run): A single trappy.Run object
                    or a path that can be passed to trappy.Run
                topology(trappy.stats.Topology): The CPU topology
                execname(str, list): List of execnames or single task
        """

        self._execnames = listify(execnames)
        self._run = Utils.init_run(run)
        self._pids = self._populate_pids()
        self._topology = topology
        self._asserts = self._populate_asserts()
        self._populate_methods()

    def _populate_asserts(self):
        """Populate SchedAsserts for the PIDs"""

        asserts = {}

        for pid in self._pids:
            asserts[pid] = SchedAssert(self._run, self._topology, pid=pid)

        return asserts

    def _populate_pids(self):
        """Map the input execnames to PIDs"""

        if len(self._execnames) == 1:
            return sconf.get_pids_for_process(self._run, self._execnames[0])

        pids = []

        for proc in self._execnames:
            pids += sconf.get_pids_for_process(self._run, proc)

        return list(set(pids))

    def _create_method(self, attr_name):
        """A wrapper function to create a dispatch function"""

        return lambda *args, **kwargs: self._dispatch(attr_name, *args, **kwargs)

    def _populate_methods(self):
        """Populate Methods from SchedAssert"""

        for attr_name in dir(SchedAssert):
            attr = getattr(SchedAssert, attr_name)

            valid_method = attr_name.startswith("get") or \
                           attr_name.startswith("assert")
            if inspect.ismethod(attr) and valid_method:
                func = self._create_method(attr_name)
                setattr(self, attr_name, func)

    def get_task_name(self, pid):
        """Get task name for the PID"""
        return self._asserts[pid].execname

    def _dispatch(self, func_name, *args, **kwargs):
        """The dispatch function to call into the SchedAssert
           Method
        """

        assert_func = func_name.startswith("assert")
        num_true = 0

        rank = kwargs.pop("rank", None)
        result = kwargs.pop("result", {})
        param = kwargs.pop("param", re.sub(r"assert|get", "", func_name, count=1).lower())

        for pid in self._pids:

            if pid not in result:
                result[pid] = {}
                result[pid]["task_name"] = self.get_task_name(pid)

            attr = getattr(self._asserts[pid], func_name)
            result[pid][param] = attr(*args, **kwargs)

            if assert_func and result[pid][param]:
                num_true += 1

        if assert_func and rank:
            return num_true == rank
        else:
            return result

    def generate_events(self, level, window=None):
        """Generate Events for the trace plot"""

        events = {}
        for s_assert in self._asserts.values():
            events[s_assert.name] = s_assert.generate_events(level, window=window)

        return events

    def plot(self, level="cpu", window=None, xlim=None):
        """
        Returns:
            trappy.plotter.AbstractDataPlotter. Call .view() for
                displaying the plot
        """

        if not xlim:
            if not window:
                xlim = [0, self._run.get_duration()]
            else:
                xlim = list(window)

        events = self.generate_events(level, window)
        names = [s.name for s in self._asserts.values()]
        num_lanes = self._topology.level_span(level)
        lane_prefix = level.upper() + ": "
        return trappy.EventPlot(events, names, lane_prefix, num_lanes, xlim)
<commit_msg>doc: SchedMultiAssert: Use sphinx friendly docstrings<commit_after>#    Copyright 2015-2015 ARM Limited
#
# 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.
#

"""A library for asserting scheduler scenarios based on the
statistics aggregation framework"""

import re
import inspect
import trappy
from trappy.stats import SchedConf as sconf
from trappy.plotter.Utils import listify
from bart.sched.SchedAssert import SchedAssert
from bart.common import Utils

class SchedMultiAssert(object):
    """This is vector assertion class built on top of
    :mod:`bart.sched.SchedAssert.SchedAssert`

    :param run: A single trappy.Run object
        or a path that can be passed to trappy.Run
    :type run: :mod:`trappy.run.Run`

    :param topology: A topology that describes the arrangement of
        CPU's on a system. This is useful for multi-cluster systems
        where data needs to be aggregated at different topological
        levels
    :type topology: :mod:`trappy.stats.Topology.Topology`

    :param execnames: The execnames of the task to be analysed

        A single execname or a list of execnames can be passed.
        There can be multiple processes associated with a single
        execname parameter. The execnames are searched using a prefix
        match.
    :type execname: list, str

    :param pid: The process ID of the task to be analysed
    :type pid: list, int

    Consider the following processes which need to be analysed

        ===== ==============
         PID    execname
        ===== ==============
         11     task_1
         22     task_2
         33     task_3
        ===== ==============

    A :mod:`bart.sched.SchedMultiAssert.SchedMultiAssert` instance be created
    following different ways:

        - Using execname prefix match
          ::

            SchedMultiAssert(run, topology, execnames="task_")

        - Individual Task names
          ::

            SchedMultiAssert(run, topology, execnames=["task_1", "task_2", "task_3"])

        - Using Process IDs
          ::

            SchedMultiAssert(run, topology, pids=[11, 22, 33])

    All the functionality provided in :mod:`bart.sched.SchedAssert.SchedAssert` is available
    in this class with the addition of handling vector assertions.

    For example consider the use of :func:`getDutyCycle`
    ::

        >>> s = SchedMultiAssert(run, topology, execnames="task_")
        >>> s.getDutyCycle(window=(start, end))
        {
            "11": {
                "task_name": "task_1",
                "dutycycle": 10.0
            },
            "22": {
                "task_name": "task_2",
                "dutycycle": 20.0
            },
            "33": {
                "task_name": "task_3",
                "dutycycle": 30.0
            },
        }

    The assertions can be used in a similar way
    ::

        >>> import operator as op
        >>> s = SchedMultiAssert(run, topology, execnames="task_")
        >>> s.assertDutyCycle(15, op.ge, window=(start, end))
        {
            "11": {
                "task_name": "task_1",
                "dutycycle": False
            },
            "22": {
                "task_name": "task_2",
                "dutycycle": True
            },
            "33": {
                "task_name": "task_3",
                "dutycycle": True
            },
        }

    The above result can be coalesced using a :code:`rank` parameter
    As we know that only 2 processes have duty cycles greater than 15%
    we can do the following:
    ::

        >>> import operator as op
        >>> s = SchedMultiAssert(run, topology, execnames="task_")
        >>> s.assertDutyCycle(15, op.ge, window=(start, end), rank=2)
        True

    See :mod:`bart.sched.SchedAssert.SchedAssert` for the available
    functionality
    """

    def __init__(self, run, topology, execnames):

        self._execnames = listify(execnames)
        self._run = Utils.init_run(run)
        self._pids = self._populate_pids()
        self._topology = topology
        self._asserts = self._populate_asserts()
        self._populate_methods()

    def _populate_asserts(self):
        """Populate SchedAsserts for the PIDs"""

        asserts = {}

        for pid in self._pids:
            asserts[pid] = SchedAssert(self._run, self._topology, pid=pid)

        return asserts

    def _populate_pids(self):
        """Map the input execnames to PIDs"""

        if len(self._execnames) == 1:
            return sconf.get_pids_for_process(self._run, self._execnames[0])

        pids = []

        for proc in self._execnames:
            pids += sconf.get_pids_for_process(self._run, proc)

        return list(set(pids))

    def _create_method(self, attr_name):
        """A wrapper function to create a dispatch function"""

        return lambda *args, **kwargs: self._dispatch(attr_name, *args, **kwargs)

    def _populate_methods(self):
        """Populate Methods from SchedAssert"""

        for attr_name in dir(SchedAssert):
            attr = getattr(SchedAssert, attr_name)

            valid_method = attr_name.startswith("get") or \
                           attr_name.startswith("assert")
            if inspect.ismethod(attr) and valid_method:
                func = self._create_method(attr_name)
                setattr(self, attr_name, func)

    def get_task_name(self, pid):
        """Get task name for the PID"""
        return self._asserts[pid].execname

    def _dispatch(self, func_name, *args, **kwargs):
        """The dispatch function to call into the SchedAssert
           Method
        """

        assert_func = func_name.startswith("assert")
        num_true = 0

        rank = kwargs.pop("rank", None)
        result = kwargs.pop("result", {})
        param = kwargs.pop("param", re.sub(r"assert|get", "", func_name, count=1).lower())

        for pid in self._pids:

            if pid not in result:
                result[pid] = {}
                result[pid]["task_name"] = self.get_task_name(pid)

            attr = getattr(self._asserts[pid], func_name)
            result[pid][param] = attr(*args, **kwargs)

            if assert_func and result[pid][param]:
                num_true += 1

        if assert_func and rank:
            return num_true == rank
        else:
            return result

    def generate_events(self, level, window=None):
        """Generate Events for the trace plot

        .. note::
            This is an internal function for plotting data
        """

        events = {}
        for s_assert in self._asserts.values():
            events[s_assert.name] = s_assert.generate_events(level, window=window)

        return events

    def plot(self, level="cpu", window=None, xlim=None):
        """
        :return: :mod:`trappy.plotter.AbstractDataPlotter` instance
            Call :func:`view` to draw the graph
        """

        if not xlim:
            if not window:
                xlim = [0, self._run.get_duration()]
            else:
                xlim = list(window)

        events = self.generate_events(level, window)
        names = [s.name for s in self._asserts.values()]
        num_lanes = self._topology.level_span(level)
        lane_prefix = level.upper() + ": "
        return trappy.EventPlot(events, names, lane_prefix, num_lanes, xlim)
<|endoftext|>


<commit_before>import xformer

from tak.model import encoding

from xformer import data, train, model
from xformer.train import hooks

from attrs import define

import torch
from torch import nn
from torch.nn import functional as F
import argparse

import typing as T  # noqa

@define
class PositionBatch:
    data: dict[str, torch.Tensor]

    @property
    def inputs(self):
        return self.data["positions"][:, :-1]

    @property
    def targets(self):
        return self.data["positions"][:, 1:]

    @property
    def mask(self):
        return self.data["mask"][:, :-1]

class MaskedARLoss:
    def __init__(self):
        self.xent = torch.nn.CrossEntropyLoss(reduction="none")

    def train_and_metrics(self, batch, logits):
        return (
            (self.xent(logits.permute(0, 2, 1), batch.targets) * batch.mask).mean(),
            {},
        )

OUTPUT_SENTINEL = 256

@define
class PositionValuePolicyBatch:
    data: dict[str, torch.Tensor]

    @property
    def inputs(self):
        return F.pad(self.data["positions"], (0, 1), value=OUTPUT_SENTINEL)

    @property
    def moves(self):
        return self.data["moves"]

    @property
    def moves_mask(self):
        return self.data["moves_mask"]

    @property
    def values(self):
        return self.data["values"]

class PolicyValueLoss:
    v_weight: float = 1.0
    policy_weight: float = 1.0

    def __init__(self):
        self.xent = torch.nn.CrossEntropyLoss(reduction="none")

    def loss_and_metrics(self, batch, logits):
        v_logits = logits["values"]
        m_logits = logits["moves"]

        # breakpoint()

        with torch.no_grad():
            argmax = torch.max(m_logits, dim=-1).indices
            match = torch.where(
                batch.moves_mask != 0,
                (argmax == batch.moves),
                torch.ones_like(argmax, dtype=torch.bool),
            )
            all_match = torch.prod(match, dim=-1).float().mean()

        v_error = F.mse_loss(v_logits, batch.values)

        return (
            self.v_weight * v_error
            + self.policy_weight
            * (
                self.xent(m_logits.permute(0, 2, 1), batch.moves) * batch.moves_mask
            ).mean()
        ), {
            "v_error": v_error.item(),
            "acc@1": all_match.mean().item(),
        }

class PolicyValueHead(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()
        self.final_ln = nn.LayerNorm(
            normalized_shape=(cfg.d_model,), dtype=dtype, device=device
        )
        self.v_proj = nn.Linear(cfg.d_model, 1, dtype=dtype, device=device)
        self.move_proj = nn.Linear(
            cfg.d_model, 3 * encoding.MAX_SLIDES, dtype=dtype, device=device
        )

    def init_weights(self, cfg):
        pass

    def forward(self, acts):
        acts = self.final_ln(acts)[:, -1]

        v = torch.tanh(self.v_proj(acts))

        moves = self.move_proj(acts).reshape(-1, 3, encoding.MAX_SLIDES)

        return {
            "values": v.squeeze(-1),
            "moves": moves,
        }

@define
class LRSchedule:
    warmup_steps: int
    cooldown_steps: int
    cooldown_start: int

    def __call__(self, stats):
        if stats.step < self.warmup_steps:
            return stats.step / self.warmup_steps
        if stats.step > self.cooldown_start:
            end = self.cooldown_start + self.cooldown_steps
            remaining = end - stats.step
            return (remaining + 1) / self.cooldown_steps
        return 1.0

def parse_args():
    parser = argparse.ArgumentParser(description="Train a transformer")
    parser.add_argument("--layers", type=int, default=2, help="Number of layers")
    parser.add_argument("--d_model", type=int, default=None, help="embedding dimension")
    parser.add_argument("--d_head", type=int, default=32, help="head dimension")
    parser.add_argument(
        "--n_ctx", type=int, default=1024, help="maximum context length"
    )
    parser.add_argument(
        "--pe", type=str, default=None, help="positional encoding (sin, learned, none)"
    )

    parser.add_argument("--data", type=str, default="data/positions", help="datasource")

    parser.add_argument("--batch", type=int, default=4, help="batch size")

    parser.add_argument(
        "--test-batches", type=int, default=16, help="number of test batches"
    )
    parser.add_argument(
        "--test-freq", type=int, default=100, help="measure test loss every N steps"
    )

    parser.add_argument(
        "--device", type=str, choices=("cpu", "cuda"), default="cuda", help="device"
    )

    parser.add_argument("--job-name", type=str, default=None, help="job name for wandb")
    parser.add_argument("--group", type=str, default=None, help="wandb group name")
    parser.add_argument("--wandb", action="store_true", default=False)
    parser.add_argument("--no-wandb", action="store_false", dest="wandb")

    parser.add_argument("--lr", type=float, default=5e-4, help="learning rate")
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--profile-steps", type=str, default=None)
    parser.add_argument("--positions", type=int, default=None)

    return parser.parse_args()

def main():
    args = parse_args()

    cfg = xformer.Config(
        n_layer=args.layers,
        d_model=args.d_model or 128 * args.layers,
        d_head=args.d_head,
        n_ctx=args.n_ctx,
        n_vocab=257,
        autoregressive_mask=False,
        output_head=PolicyValueHead,
    )
    if args.pe is not None:
        cfg.positional_encoding = args.pe

    train_ds = data.Dataset(
        args.data + "-train.pt",
        batch_size=args.batch,
        device=args.device,
        batch_class=PositionValuePolicyBatch,
    )
    test_ds = data.Dataset(
        args.data + "-test.pt",
        batch_size=args.batch,
        device=args.device,
        batches=args.test_batches,
        batch_class=PositionValuePolicyBatch,
    )

    extra_hooks = []
    if args.wandb:
        extra_hooks.append(
            hooks.Wandb(
                project="taktician",
                job_name=args.job_name,
                group=args.group,
                config=args,
            )
        )
    if args.profile_steps:
        extra_hooks.append(
            hooks.Profile(
                extra_steps=set(map(int, args.profile_steps.split(","))),
            )
        )

    if args.steps:
        warmup_frac = 0.05
        cooldown_frac = 0.8
        schedule = LRSchedule(
            warmup_steps=int(warmup_frac * args.steps),
            cooldown_start=int((1 - cooldown_frac) * args.steps),
            cooldown_steps=int(cooldown_frac * args.steps),
        )
    else:
        schedule = None

    run = train.Run(
        model=model.Transformer(cfg, dtype=torch.float32, device=args.device),
        dataset=train_ds,
        # loss=MaskedARLoss(),
        loss=PolicyValueLoss(),
        optimizer=train.Optimizer(lr=args.lr, lr_schedule=schedule),
        stop=train.StopTrigger(steps=args.steps, sequences=args.positions),
        hooks=[
            hooks.TestLoss(test_ds, args.test_freq),
        ]
        + extra_hooks,
    )

    print(
        f"Training a {cfg.n_layer}L model with {cfg.n_parameters:,} non-embedding parameters..."
    )
    param_bytes = sum(t.numel() * t.element_size() for t in run.model.parameters())
    print(f" Model params use {param_bytes/1024**3:.2f}GiB on device")

    trainer = train.Trainer(run)
    trainer.train()

if __name__ == "__main__":
    main()
<commit_msg>SaveHook<commit_after>import xformer

from tak.model import encoding

from xformer import data, train, model
from xformer.train import hooks

from attrs import define

import torch
from torch import nn
from torch.nn import functional as F
import argparse

import typing as T  # noqa

@define
class PositionBatch:
    data: dict[str, torch.Tensor]

    @property
    def inputs(self):
        return self.data["positions"][:, :-1]

    @property
    def targets(self):
        return self.data["positions"][:, 1:]

    @property
    def mask(self):
        return self.data["mask"][:, :-1]

class MaskedARLoss:
    def __init__(self):
        self.xent = torch.nn.CrossEntropyLoss(reduction="none")

    def train_and_metrics(self, batch, logits):
        return (
            (self.xent(logits.permute(0, 2, 1), batch.targets) * batch.mask).mean(),
            {},
        )

OUTPUT_SENTINEL = 256

@define
class PositionValuePolicyBatch:
    data: dict[str, torch.Tensor]

    @property
    def inputs(self):
        return F.pad(self.data["positions"], (0, 1), value=OUTPUT_SENTINEL)

    @property
    def moves(self):
        return self.data["moves"]

    @property
    def moves_mask(self):
        return self.data["moves_mask"]

    @property
    def values(self):
        return self.data["values"]

class PolicyValueLoss:
    v_weight: float = 1.0
    policy_weight: float = 1.0

    def __init__(self):
        self.xent = torch.nn.CrossEntropyLoss(reduction="none")

    def loss_and_metrics(self, batch, logits):
        v_logits = logits["values"]
        m_logits = logits["moves"]

        # breakpoint()

        with torch.no_grad():
            argmax = torch.max(m_logits, dim=-1).indices
            match = torch.where(
                batch.moves_mask != 0,
                (argmax == batch.moves),
                torch.ones_like(argmax, dtype=torch.bool),
            )
            all_match = torch.prod(match, dim=-1).float().mean()

        v_error = F.mse_loss(v_logits, batch.values)

        return (
            self.v_weight * v_error
            + self.policy_weight
            * (
                self.xent(m_logits.permute(0, 2, 1), batch.moves) * batch.moves_mask
            ).mean()
        ), {
            "v_error": v_error.item(),
            "acc@1": all_match.mean().item(),
        }

class PolicyValueHead(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()
        self.final_ln = nn.LayerNorm(
            normalized_shape=(cfg.d_model,), dtype=dtype, device=device
        )
        self.v_proj = nn.Linear(cfg.d_model, 1, dtype=dtype, device=device)
        self.move_proj = nn.Linear(
            cfg.d_model, 3 * encoding.MAX_SLIDES, dtype=dtype, device=device
        )

    def init_weights(self, cfg):
        pass

    def forward(self, acts):
        acts = self.final_ln(acts)[:, -1]

        v = torch.tanh(self.v_proj(acts))

        moves = self.move_proj(acts).reshape(-1, 3, encoding.MAX_SLIDES)

        return {
            "values": v.squeeze(-1),
            "moves": moves,
        }

@define
class LRSchedule:
    warmup_steps: int
    cooldown_steps: int
    cooldown_start: int

    def __call__(self, stats):
        if stats.step < self.warmup_steps:
            return stats.step / self.warmup_steps
        if stats.step > self.cooldown_start:
            end = self.cooldown_start + self.cooldown_steps
            remaining = end - stats.step
            return (remaining + 1) / self.cooldown_steps
        return 1.0

@define
class SaveHook(train.Hook):
    save_dir: str
    step_freq: int

    def after_step(self, run, stats):
        if stats.step % self.step_freq != 0:
            return
        self.save_run(run, stats)

    def save_run(self, run: train.Run, stats: train.Stats):
        run_dir = os.path.join(self.save_dir, f"step_{stats.step:06d}")
        os.makedirs(run_dir, exist_ok=True)
        print(f"Saving to {run_dir}...")
        torch.save(
            run.model.state_dict(),
            os.path.join(run_dir, "model.pt"),
        )
        with open(os.path.join(run_dir, "config.yaml"), "w") as fh:
            yaml.dump(run.model.cfg, fh)
        with open(os.path.join(run_dir, "stats.json"), "w") as fh:
            json.dump(attrs.asdict(stats), fh, indent=2)
        # torch.save(os.path.join(run_dir, "model.opt.pt"), run.optimizer.state_dict())

def parse_args():
    parser = argparse.ArgumentParser(description="Train a transformer")
    parser.add_argument("--layers", type=int, default=2, help="Number of layers")
    parser.add_argument("--d_model", type=int, default=None, help="embedding dimension")
    parser.add_argument("--d_head", type=int, default=32, help="head dimension")
    parser.add_argument(
        "--n_ctx", type=int, default=1024, help="maximum context length"
    )
    parser.add_argument(
        "--pe", type=str, default=None, help="positional encoding (sin, learned, none)"
    )

    parser.add_argument("--data", type=str, default="data/positions", help="datasource")

    parser.add_argument("--batch", type=int, default=4, help="batch size")

    parser.add_argument(
        "--test-batches", type=int, default=16, help="number of test batches"
    )
    parser.add_argument(
        "--test-freq", type=int, default=100, help="measure test loss every N steps"
    )

    parser.add_argument(
        "--save-freq", type=int, default=100, help="save model every N steps"
    )
    parser.add_argument("--save-dir", type=str, help="save directory")

    parser.add_argument(
        "--device", type=str, choices=("cpu", "cuda"), default="cuda", help="device"
    )

    parser.add_argument("--job-name", type=str, default=None, help="job name for wandb")
    parser.add_argument("--group", type=str, default=None, help="wandb group name")
    parser.add_argument("--wandb", action="store_true", default=False)
    parser.add_argument("--no-wandb", action="store_false", dest="wandb")

    parser.add_argument("--lr", type=float, default=5e-4, help="learning rate")
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--profile-steps", type=str, default=None)
    parser.add_argument("--positions", type=int, default=None)

    return parser.parse_args()

def main():
    args = parse_args()

    cfg = xformer.Config(
        n_layer=args.layers,
        d_model=args.d_model or 128 * args.layers,
        d_head=args.d_head,
        n_ctx=args.n_ctx,
        n_vocab=257,
        autoregressive_mask=False,
        output_head=PolicyValueHead,
    )
    if args.pe is not None:
        cfg.positional_encoding = args.pe

    train_ds = data.Dataset(
        args.data + "-train.pt",
        batch_size=args.batch,
        device=args.device,
        batch_class=PositionValuePolicyBatch,
    )
    test_ds = data.Dataset(
        args.data + "-test.pt",
        batch_size=args.batch,
        device=args.device,
        batches=args.test_batches,
        batch_class=PositionValuePolicyBatch,
    )

    extra_hooks = []
    if args.wandb:
        extra_hooks.append(
            hooks.Wandb(
                project="taktician",
                job_name=args.job_name,
                group=args.group,
                config=args,
            )
        )
    if args.profile_steps:
        extra_hooks.append(
            hooks.Profile(
                extra_steps=set(map(int, args.profile_steps.split(","))),
            )
        )
    if args.save_dir:
        extra_hooks.append(
            SaveHook(
                save_dir=args.save_dir,
                step_freq=args.save_freq,
            )
        )

    if args.steps:
        warmup_frac = 0.05
        cooldown_frac = 0.8
        schedule = LRSchedule(
            warmup_steps=int(warmup_frac * args.steps),
            cooldown_start=int((1 - cooldown_frac) * args.steps),
            cooldown_steps=int(cooldown_frac * args.steps),
        )
    else:
        schedule = None

    run = train.Run(
        model=model.Transformer(cfg, dtype=torch.float32, device=args.device),
        dataset=train_ds,
        # loss=MaskedARLoss(),
        loss=PolicyValueLoss(),
        optimizer=train.Optimizer(lr=args.lr, lr_schedule=schedule),
        stop=train.StopTrigger(steps=args.steps, sequences=args.positions),
        hooks=[
            hooks.TestLoss(test_ds, args.test_freq),
        ]
        + extra_hooks,
    )

    print(
        f"Training a {cfg.n_layer}L model with {cfg.n_parameters:,} non-embedding parameters..."
    )
    param_bytes = sum(t.numel() * t.element_size() for t in run.model.parameters())
    print(f" Model params use {param_bytes/1024**3:.2f}GiB on device")

    trainer = train.Trainer(run)
    trainer.train()

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


<commit_before>import sys
import unittest
import weakref

from test_support import run_unittest

class C:
    def method(self):
        pass

class Callable:
    bar = None

    def __call__(self, x):
        self.bar = x

def create_function():
    def f(): pass
    return f

def create_bound_method():
    return C().method

def create_unbound_method():
    return C.method

class TestBase(unittest.TestCase):

    def setUp(self):
        self.cbcalled = 0

    def callback(self, ref):
        self.cbcalled += 1

class ReferencesTestCase(TestBase):

    def test_basic_ref(self):
        self.check_basic_ref(C)
        self.check_basic_ref(create_function)
        self.check_basic_ref(create_bound_method)
        self.check_basic_ref(create_unbound_method)

    def test_basic_callback(self):
        self.check_basic_callback(C)
        self.check_basic_callback(create_function)
        self.check_basic_callback(create_bound_method)
        self.check_basic_callback(create_unbound_method)

    def test_multiple_callbacks(self):
        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del o
        self.assert_(ref1() is None,
                     "expected reference to be invalidated")
        self.assert_(ref2() is None,
                     "expected reference to be invalidated")
        self.assert_(self.cbcalled == 2,
                     "callback not called the right number of times")

    def test_multiple_selfref_callbacks(self):
        """Make sure all references are invalidated before callbacks
        are called."""
        #
        # What's important here is that we're using the first
        # reference in the callback invoked on the second reference
        # (the most recently created ref is cleaned up first).  This
        # tests that all references to the object are invalidated
        # before any of the callbacks are invoked, so that we only
        # have one invocation of _weakref.c:cleanup_helper() active
        # for a particular object at a time.
        #
        def callback(object, self=self):
            self.ref()
        c = C()
        self.ref = weakref.ref(c, callback)
        ref1 = weakref.ref(c, callback)
        del c

    def test_proxy_ref(self):
        o = C()
        o.bar = 1
        ref1 = weakref.proxy(o, self.callback)
        ref2 = weakref.proxy(o, self.callback)
        del o

        def check(proxy):
            proxy.bar

        self.assertRaises(weakref.ReferenceError, check, ref1)
        self.assertRaises(weakref.ReferenceError, check, ref2)
        self.assert_(self.cbcalled == 2)

    def check_basic_ref(self, factory):
        o = factory()
        ref = weakref.ref(o)
        self.assert_(ref() is not None,
                     "weak reference to live object should be live")
        o2 = ref()
        self.assert_(o is o2,
                     "<ref>() should return original object if live")

    def check_basic_callback(self, factory):
        self.cbcalled = 0
        o = factory()
        ref = weakref.ref(o, self.callback)
        del o
        self.assert_(self.cbcalled == 1,
                     "callback did not properly set 'cbcalled'")
        self.assert_(ref() is None,
                     "ref2 should be dead after deleting object reference")

    def test_ref_reuse(self):
        o = C()
        ref1 = weakref.ref(o)
        # create a proxy to make sure that there's an intervening creation
        # between these two; it should make no difference
        proxy = weakref.proxy(o)
        ref2 = weakref.ref(o)
        self.assert_(ref1 is ref2,
                     "reference object w/out callback should be re-used")

        o = C()
        proxy = weakref.proxy(o)
        ref1 = weakref.ref(o)
        ref2 = weakref.ref(o)
        self.assert_(ref1 is ref2,
                     "reference object w/out callback should be re-used")
        self.assert_(weakref.getweakrefcount(o) == 2,
                     "wrong weak ref count for object")
        del proxy
        self.assert_(weakref.getweakrefcount(o) == 1,
                     "wrong weak ref count for object after deleting proxy")

    def test_proxy_reuse(self):
        o = C()
        proxy1 = weakref.proxy(o)
        ref = weakref.ref(o)
        proxy2 = weakref.proxy(o)
        self.assert_(proxy1 is proxy2,
                     "proxy object w/out callback should have been re-used")

    def test_basic_proxy(self):
        o = C()
        self.check_proxy(o, weakref.proxy(o))

    def test_callable_proxy(self):
        o = Callable()
        ref1 = weakref.proxy(o)

        self.check_proxy(o, ref1)

        self.assert_(type(ref1) is weakref.CallableProxyType,
                     "proxy is not of callable type")
        ref1('twinkies!')
        self.assert_(o.bar == 'twinkies!',
                     "call through proxy not passed through to original")

        # expect due to too few args
        self.assertRaises(TypeError, ref1)

        # expect due to too many args
        self.assertRaises(TypeError, ref1, 1, 2, 3)

    def check_proxy(self, o, proxy):
        o.foo = 1
        self.assert_(proxy.foo == 1,
                     "proxy does not reflect attribute addition")
        o.foo = 2
        self.assert_(proxy.foo == 2,
                     "proxy does not reflect attribute modification")
        del o.foo
        self.assert_(not hasattr(proxy, 'foo'),
                     "proxy does not reflect attribute removal")

        proxy.foo = 1
        self.assert_(o.foo == 1,
                     "object does not reflect attribute addition via proxy")
        proxy.foo = 2
        self.assert_(
            o.foo == 2,
            "object does not reflect attribute modification via proxy")
        del proxy.foo
        self.assert_(not hasattr(o, 'foo'),
                     "object does not reflect attribute removal via proxy")

    def test_getweakrefcount(self):
        o = C()
        ref1 = weakref.ref(o)
        ref2 = weakref.ref(o, self.callback)
        self.assert_(weakref.getweakrefcount(o) == 2,
                     "got wrong number of weak reference objects")

        proxy1 = weakref.proxy(o)
        proxy2 = weakref.proxy(o, self.callback)
        self.assert_(weakref.getweakrefcount(o) == 4,
                     "got wrong number of weak reference objects")

    def test_getweakrefs(self):
        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del ref1
        self.assert_(weakref.getweakrefs(o) == [ref2],
                     "list of refs does not match")

        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del ref2
        self.assert_(weakref.getweakrefs(o) == [ref1],
                     "list of refs does not match")

class Object:
    def __init__(self, arg):
        self.arg = arg
    def __repr__(self):
        return "<Object %r>" % self.arg

class MappingTestCase(TestBase):

    COUNT = 10

    def test_weak_values(self):
        #
        #  This exercises d.copy(), d.items(), d[], del d[], len(d).
        #
        dict, objects = self.make_weak_valued_dict()
        for o in objects:
            self.assert_(weakref.getweakrefcount(o) == 1,
                         "wrong number of weak references to %r!" % o)
            self.assert_(o is dict[o.arg],
                         "wrong object returned by weak dict!")
        items1 = dict.items()
        items2 = dict.copy().items()
        items1.sort()
        items2.sort()
        self.assert_(items1 == items2,
                     "cloning of weak-valued dictionary did not work!")
        del items1, items2
        self.assert_(len(dict) == self.COUNT)
        del objects[0]
        self.assert_(len(dict) == (self.COUNT - 1),
                     "deleting object did not cause dictionary update")
        del objects, o
        self.assert_(len(dict) == 0,
                     "deleting the values did not clear the dictionary")

    def test_weak_keys(self):
        #
        #  This exercises d.copy(), d.items(), d[] = v, d[], del d[],
        #  len(d).
        #
        dict, objects = self.make_weak_keyed_dict()
        for o in objects:
            self.assert_(weakref.getweakrefcount(o) == 1,
                         "wrong number of weak references to %r!" % o)
            self.assert_(o.arg is dict[o],
                         "wrong object returned by weak dict!")
        items1 = dict.items()
        items2 = dict.copy().items()
        items1.sort()
        items2.sort()
        self.assert_(items1 == items2,
                     "cloning of weak-keyed dictionary did not work!")
        del items1, items2
        self.assert_(len(dict) == self.COUNT)
        del objects[0]
        self.assert_(len(dict) == (self.COUNT - 1),
                     "deleting object did not cause dictionary update")
        del objects, o
        self.assert_(len(dict) == 0,
                     "deleting the keys did not clear the dictionary")

    def test_weak_keyed_iters(self):
        dict, objects = self.make_weak_keyed_dict()
        self.check_iters(dict)

    def test_weak_valued_iters(self):
        dict, objects = self.make_weak_valued_dict()
        self.check_iters(dict)

    def check_iters(self, dict):
        # item iterator:
        items = dict.items()
        for item in dict.iteritems():
            items.remove(item)
        self.assert_(len(items) == 0, "iterator did not touch all items")

        # key iterator:
        keys = dict.keys()
        for k in dict:
            keys.remove(k)
        self.assert_(len(keys) == 0, "iterator did not touch all keys")

        # value iterator:
        values = dict.values()
        for v in dict.itervalues():
            values.remove(v)
        self.assert_(len(values) == 0, "iterator did not touch all values")

    def make_weak_keyed_dict(self):
        dict = weakref.WeakKeyDictionary()
        objects = map(Object, range(self.COUNT))
        for o in objects:
            dict[o] = o.arg
        return dict, objects

    def make_weak_valued_dict(self):
        dict = weakref.WeakValueDictionary()
        objects = map(Object, range(self.COUNT))
        for o in objects:
            dict[o.arg] = o
        return dict, objects

    def check_update(self, klass, dict):
        #
        #  This exercises d.update(), len(d), d.keys(), d.has_key(),
        #  d.get(), d[].
        #
        weakdict = klass()
        weakdict.update(dict)
        self.assert_(len(weakdict) == len(dict))
        for k in weakdict.keys():
            self.assert_(dict.has_key(k),
                         "mysterious new key appeared in weak dict")
            v = dict.get(k)
            self.assert_(v is weakdict[k])
            self.assert_(v is weakdict.get(k))
        for k in dict.keys():
            self.assert_(weakdict.has_key(k),
                         "original key disappeared in weak dict")
            v = dict[k]
            self.assert_(v is weakdict[k])
            self.assert_(v is weakdict.get(k))

    def test_weak_valued_dict_update(self):
        self.check_update(weakref.WeakValueDictionary,
                          {1: C(), 'a': C(), C(): C()})

    def test_weak_keyed_dict_update(self):
        self.check_update(weakref.WeakKeyDictionary,
                          {C(): 1, C(): 2, C(): 3})

run_unittest(ReferencesTestCase)
run_unittest(MappingTestCase)
<commit_msg>Extend the weakref test suite to cover the complete mapping interface for both weakref.Weak*Dictionary classes.<commit_after>import sys
import unittest
import weakref

from test_support import run_unittest

class C:
    def method(self):
        pass

class Callable:
    bar = None

    def __call__(self, x):
        self.bar = x

def create_function():
    def f(): pass
    return f

def create_bound_method():
    return C().method

def create_unbound_method():
    return C.method

class TestBase(unittest.TestCase):

    def setUp(self):
        self.cbcalled = 0

    def callback(self, ref):
        self.cbcalled += 1

class ReferencesTestCase(TestBase):

    def test_basic_ref(self):
        self.check_basic_ref(C)
        self.check_basic_ref(create_function)
        self.check_basic_ref(create_bound_method)
        self.check_basic_ref(create_unbound_method)

    def test_basic_callback(self):
        self.check_basic_callback(C)
        self.check_basic_callback(create_function)
        self.check_basic_callback(create_bound_method)
        self.check_basic_callback(create_unbound_method)

    def test_multiple_callbacks(self):
        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del o
        self.assert_(ref1() is None,
                     "expected reference to be invalidated")
        self.assert_(ref2() is None,
                     "expected reference to be invalidated")
        self.assert_(self.cbcalled == 2,
                     "callback not called the right number of times")

    def test_multiple_selfref_callbacks(self):
        """Make sure all references are invalidated before callbacks
        are called."""
        #
        # What's important here is that we're using the first
        # reference in the callback invoked on the second reference
        # (the most recently created ref is cleaned up first).  This
        # tests that all references to the object are invalidated
        # before any of the callbacks are invoked, so that we only
        # have one invocation of _weakref.c:cleanup_helper() active
        # for a particular object at a time.
        #
        def callback(object, self=self):
            self.ref()
        c = C()
        self.ref = weakref.ref(c, callback)
        ref1 = weakref.ref(c, callback)
        del c

    def test_proxy_ref(self):
        o = C()
        o.bar = 1
        ref1 = weakref.proxy(o, self.callback)
        ref2 = weakref.proxy(o, self.callback)
        del o

        def check(proxy):
            proxy.bar

        self.assertRaises(weakref.ReferenceError, check, ref1)
        self.assertRaises(weakref.ReferenceError, check, ref2)
        self.assert_(self.cbcalled == 2)

    def check_basic_ref(self, factory):
        o = factory()
        ref = weakref.ref(o)
        self.assert_(ref() is not None,
                     "weak reference to live object should be live")
        o2 = ref()
        self.assert_(o is o2,
                     "<ref>() should return original object if live")

    def check_basic_callback(self, factory):
        self.cbcalled = 0
        o = factory()
        ref = weakref.ref(o, self.callback)
        del o
        self.assert_(self.cbcalled == 1,
                     "callback did not properly set 'cbcalled'")
        self.assert_(ref() is None,
                     "ref2 should be dead after deleting object reference")

    def test_ref_reuse(self):
        o = C()
        ref1 = weakref.ref(o)
        # create a proxy to make sure that there's an intervening creation
        # between these two; it should make no difference
        proxy = weakref.proxy(o)
        ref2 = weakref.ref(o)
        self.assert_(ref1 is ref2,
                     "reference object w/out callback should be re-used")

        o = C()
        proxy = weakref.proxy(o)
        ref1 = weakref.ref(o)
        ref2 = weakref.ref(o)
        self.assert_(ref1 is ref2,
                     "reference object w/out callback should be re-used")
        self.assert_(weakref.getweakrefcount(o) == 2,
                     "wrong weak ref count for object")
        del proxy
        self.assert_(weakref.getweakrefcount(o) == 1,
                     "wrong weak ref count for object after deleting proxy")

    def test_proxy_reuse(self):
        o = C()
        proxy1 = weakref.proxy(o)
        ref = weakref.ref(o)
        proxy2 = weakref.proxy(o)
        self.assert_(proxy1 is proxy2,
                     "proxy object w/out callback should have been re-used")

    def test_basic_proxy(self):
        o = C()
        self.check_proxy(o, weakref.proxy(o))

    def test_callable_proxy(self):
        o = Callable()
        ref1 = weakref.proxy(o)

        self.check_proxy(o, ref1)

        self.assert_(type(ref1) is weakref.CallableProxyType,
                     "proxy is not of callable type")
        ref1('twinkies!')
        self.assert_(o.bar == 'twinkies!',
                     "call through proxy not passed through to original")

        # expect due to too few args
        self.assertRaises(TypeError, ref1)

        # expect due to too many args
        self.assertRaises(TypeError, ref1, 1, 2, 3)

    def check_proxy(self, o, proxy):
        o.foo = 1
        self.assert_(proxy.foo == 1,
                     "proxy does not reflect attribute addition")
        o.foo = 2
        self.assert_(proxy.foo == 2,
                     "proxy does not reflect attribute modification")
        del o.foo
        self.assert_(not hasattr(proxy, 'foo'),
                     "proxy does not reflect attribute removal")

        proxy.foo = 1
        self.assert_(o.foo == 1,
                     "object does not reflect attribute addition via proxy")
        proxy.foo = 2
        self.assert_(
            o.foo == 2,
            "object does not reflect attribute modification via proxy")
        del proxy.foo
        self.assert_(not hasattr(o, 'foo'),
                     "object does not reflect attribute removal via proxy")

    def test_getweakrefcount(self):
        o = C()
        ref1 = weakref.ref(o)
        ref2 = weakref.ref(o, self.callback)
        self.assert_(weakref.getweakrefcount(o) == 2,
                     "got wrong number of weak reference objects")

        proxy1 = weakref.proxy(o)
        proxy2 = weakref.proxy(o, self.callback)
        self.assert_(weakref.getweakrefcount(o) == 4,
                     "got wrong number of weak reference objects")

    def test_getweakrefs(self):
        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del ref1
        self.assert_(weakref.getweakrefs(o) == [ref2],
                     "list of refs does not match")

        o = C()
        ref1 = weakref.ref(o, self.callback)
        ref2 = weakref.ref(o, self.callback)
        del ref2
        self.assert_(weakref.getweakrefs(o) == [ref1],
                     "list of refs does not match")

class Object:
    def __init__(self, arg):
        self.arg = arg
    def __repr__(self):
        return "<Object %r>" % self.arg

class MappingTestCase(TestBase):

    COUNT = 10

    def test_weak_values(self):
        #
        #  This exercises d.copy(), d.items(), d[], del d[], len(d).
        #
        dict, objects = self.make_weak_valued_dict()
        for o in objects:
            self.assert_(weakref.getweakrefcount(o) == 1,
                         "wrong number of weak references to %r!" % o)
            self.assert_(o is dict[o.arg],
                         "wrong object returned by weak dict!")
        items1 = dict.items()
        items2 = dict.copy().items()
        items1.sort()
        items2.sort()
        self.assert_(items1 == items2,
                     "cloning of weak-valued dictionary did not work!")
        del items1, items2
        self.assert_(len(dict) == self.COUNT)
        del objects[0]
        self.assert_(len(dict) == (self.COUNT - 1),
                     "deleting object did not cause dictionary update")
        del objects, o
        self.assert_(len(dict) == 0,
                     "deleting the values did not clear the dictionary")

    def test_weak_keys(self):
        #
        #  This exercises d.copy(), d.items(), d[] = v, d[], del d[],
        #  len(d).
        #
        dict, objects = self.make_weak_keyed_dict()
        for o in objects:
            self.assert_(weakref.getweakrefcount(o) == 1,
                         "wrong number of weak references to %r!" % o)
            self.assert_(o.arg is dict[o],
                         "wrong object returned by weak dict!")
        items1 = dict.items()
        items2 = dict.copy().items()
        items1.sort()
        items2.sort()
        self.assert_(items1 == items2,
                     "cloning of weak-keyed dictionary did not work!")
        del items1, items2
        self.assert_(len(dict) == self.COUNT)
        del objects[0]
        self.assert_(len(dict) == (self.COUNT - 1),
                     "deleting object did not cause dictionary update")
        del objects, o
        self.assert_(len(dict) == 0,
                     "deleting the keys did not clear the dictionary")

    def test_weak_keyed_iters(self):
        dict, objects = self.make_weak_keyed_dict()
        self.check_iters(dict)

    def test_weak_valued_iters(self):
        dict, objects = self.make_weak_valued_dict()
        self.check_iters(dict)

    def check_iters(self, dict):
        # item iterator:
        items = dict.items()
        for item in dict.iteritems():
            items.remove(item)
        self.assert_(len(items) == 0, "iteritems() did not touch all items")

        # key iterator, via __iter__():
        keys = dict.keys()
        for k in dict:
            keys.remove(k)
        self.assert_(len(keys) == 0, "__iter__() did not touch all keys")

        # key iterator, via iterkeys():
        keys = dict.keys()
        for k in dict.iterkeys():
            keys.remove(k)
        self.assert_(len(keys) == 0, "iterkeys() did not touch all keys")

        # value iterator:
        values = dict.values()
        for v in dict.itervalues():
            values.remove(v)
        self.assert_(len(values) == 0, "itervalues() did not touch all values")

    def make_weak_keyed_dict(self):
        dict = weakref.WeakKeyDictionary()
        objects = map(Object, range(self.COUNT))
        for o in objects:
            dict[o] = o.arg
        return dict, objects

    def make_weak_valued_dict(self):
        dict = weakref.WeakValueDictionary()
        objects = map(Object, range(self.COUNT))
        for o in objects:
            dict[o.arg] = o
        return dict, objects

    def check_popitem(self, klass, key1, value1, key2, value2):
        weakdict = klass()
        weakdict[key1] = value1
        weakdict[key2] = value2
        self.assert_(len(weakdict) == 2)
        k, v = weakdict.popitem()
        self.assert_(len(weakdict) == 1)
        if k is key1:
            self.assert_(v is value1)
        else:
            self.assert_(v is value2)
        k, v = weakdict.popitem()
        self.assert_(len(weakdict) == 0)
        if k is key1:
            self.assert_(v is value1)
        else:
            self.assert_(v is value2)

    def test_weak_valued_dict_popitem(self):
        self.check_popitem(weakref.WeakValueDictionary,
                           "key1", C(), "key2", C())

    def test_weak_keyed_dict_popitem(self):
        self.check_popitem(weakref.WeakKeyDictionary,
                           C(), "value 1", C(), "value 2")

    def check_setdefault(self, klass, key, value1, value2):
        self.assert_(value1 is not value2,
                     "invalid test"
                     " -- value parameters must be distinct objects")
        weakdict = klass()
        o = weakdict.setdefault(key, value1)
        self.assert_(o is value1)
        self.assert_(weakdict.has_key(key))
        self.assert_(weakdict.get(key) is value1)
        self.assert_(weakdict[key] is value1)

        o = weakdict.setdefault(key, value2)
        self.assert_(o is value1)
        self.assert_(weakdict.has_key(key))
        self.assert_(weakdict.get(key) is value1)
        self.assert_(weakdict[key] is value1)

    def test_weak_valued_dict_setdefault(self):
        self.check_setdefault(weakref.WeakValueDictionary,
                              "key", C(), C())

    def test_weak_keyed_dict_setdefault(self):
        self.check_setdefault(weakref.WeakKeyDictionary,
                              C(), "value 1", "value 2")

    def check_update(self, klass, dict):
        #
        #  This exercises d.update(), len(d), d.keys(), d.has_key(),
        #  d.get(), d[].
        #
        weakdict = klass()
        weakdict.update(dict)
        self.assert_(len(weakdict) == len(dict))
        for k in weakdict.keys():
            self.assert_(dict.has_key(k),
                         "mysterious new key appeared in weak dict")
            v = dict.get(k)
            self.assert_(v is weakdict[k])
            self.assert_(v is weakdict.get(k))
        for k in dict.keys():
            self.assert_(weakdict.has_key(k),
                         "original key disappeared in weak dict")
            v = dict[k]
            self.assert_(v is weakdict[k])
            self.assert_(v is weakdict.get(k))

    def test_weak_valued_dict_update(self):
        self.check_update(weakref.WeakValueDictionary,
                          {1: C(), 'a': C(), C(): C()})

    def test_weak_keyed_dict_update(self):
        self.check_update(weakref.WeakKeyDictionary,
                          {C(): 1, C(): 2, C(): 3})

run_unittest(ReferencesTestCase)
run_unittest(MappingTestCase)
<|endoftext|>


<commit_before>"""
Tests for giving invalid JSON to endpoints.
"""

from datetime import datetime, timedelta

import pytest
import pytz
import requests
from freezegun import freeze_time
from requests import codes

from mock_vws._constants import ResultCodes
from tests.mock_vws.utils import (
    TargetAPIEndpoint,
    assert_vws_failure,
    authorization_header,
    rfc_1123_date,
)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestInvalidJSON:
    """
    Tests for giving invalid JSON to endpoints.
    """

    @pytest.mark.parametrize('date_skew_minutes', [0, 10])
    def test_invalid_json(
        self,
        endpoint: TargetAPIEndpoint,
        date_skew_minutes: int,
    ) -> None:
        """
        Giving invalid JSON to endpoints returns error responses.
        """
        date_is_skewed = not date_skew_minutes == 0
        # This is an undocumented difference between `/summary` and other
        # endpoints.
        is_summary_endpoint = endpoint.prepared_request.path_url == '/summary'
        content = b'a'
        gmt = pytz.timezone('GMT')
        now = datetime.now(tz=gmt)
        time_to_freeze = now + timedelta(minutes=date_skew_minutes)
        with freeze_time(time_to_freeze):
            date = rfc_1123_date()

        endpoint_headers = dict(endpoint.prepared_request.headers)
        authorization_string = authorization_header(
            access_key=endpoint.access_key,
            secret_key=endpoint.secret_key,
            method=str(endpoint.prepared_request.method),
            content=content,
            content_type=endpoint.auth_header_content_type,
            date=date,
            request_path=endpoint.prepared_request.path_url,
        )

        headers = {
            **endpoint_headers,
            'Authorization': authorization_string,
            'Date': date,
        }

        endpoint.prepared_request.prepare_body(  # type: ignore
            data=content,
            files=None,
        )

        endpoint.prepared_request.prepare_headers(  # type: ignore
            headers=headers,
        )
        endpoint.prepared_request.prepare_content_length(  # type: ignore
            body=content,
        )
        session = requests.Session()
        response = session.send(  # type: ignore
            request=endpoint.prepared_request,
        )

        takes_json_data = (
            endpoint.auth_header_content_type == 'application/json'
        )

        if date_is_skewed and takes_json_data:
            # On the real implementation, we get `codes.FORBIDDEN` and
            # `REQUEST_TIME_TOO_SKEWED`.
            # See https://github.com/adamtheturtle/vws-python/issues/407 for
            # implementing this on them mock.
            return

        if not date_is_skewed and takes_json_data:
            assert_vws_failure(
                response=response,
                status_code=codes.BAD_REQUEST,
                result_code=ResultCodes.FAIL,
            )
            return

        if date_is_skewed and is_summary_endpoint:
            assert_vws_failure(
                response=response,
                status_code=codes.FORBIDDEN,
                result_code=ResultCodes.REQUEST_TIME_TOO_SKEWED,
            )
            return

        if not date_is_skewed and is_summary_endpoint:
            assert_vws_failure(
                response=response,
                status_code=codes.UNAUTHORIZED,
                result_code=ResultCodes.AUTHENTICATION_FAILURE,
            )
            return

        assert response.status_code == codes.BAD_REQUEST
        assert response.text == ''
        assert 'Content-Type' not in response.headers
<commit_msg>Add test for invalid JSON given to query endpoint<commit_after>"""
Tests for giving invalid JSON to endpoints.
"""

from datetime import datetime, timedelta
from urllib.parse import urlparse

import pytest
import pytz
import requests
from freezegun import freeze_time
from requests import codes

from mock_vws._constants import ResultCodes
from tests.mock_vws.utils import (
    TargetAPIEndpoint,
    assert_vwq_failure,
    assert_vws_failure,
    authorization_header,
    rfc_1123_date,
)

@pytest.mark.usefixtures('verify_mock_vuforia')
class TestInvalidJSON:
    """
    Tests for giving invalid JSON to endpoints.
    """

    @pytest.mark.parametrize('date_skew_minutes', [0, 10])
    def test_invalid_json(
        self,
        endpoint: TargetAPIEndpoint,
        date_skew_minutes: int,
    ) -> None:
        """
        Giving invalid JSON to endpoints returns error responses.
        """
        date_is_skewed = not date_skew_minutes == 0
        # This is an undocumented difference between `/summary` and other
        # endpoints.
        is_summary_endpoint = endpoint.prepared_request.path_url == '/summary'
        content = b'a'
        gmt = pytz.timezone('GMT')
        now = datetime.now(tz=gmt)
        time_to_freeze = now + timedelta(minutes=date_skew_minutes)
        with freeze_time(time_to_freeze):
            date = rfc_1123_date()

        endpoint_headers = dict(endpoint.prepared_request.headers)
        authorization_string = authorization_header(
            access_key=endpoint.access_key,
            secret_key=endpoint.secret_key,
            method=str(endpoint.prepared_request.method),
            content=content,
            content_type=endpoint.auth_header_content_type,
            date=date,
            request_path=endpoint.prepared_request.path_url,
        )

        headers = {
            **endpoint_headers,
            'Authorization': authorization_string,
            'Date': date,
        }

        endpoint.prepared_request.prepare_body(  # type: ignore
            data=content,
            files=None,
        )

        endpoint.prepared_request.prepare_headers(  # type: ignore
            headers=headers,
        )
        endpoint.prepared_request.prepare_content_length(  # type: ignore
            body=content,
        )
        session = requests.Session()
        response = session.send(  # type: ignore
            request=endpoint.prepared_request,
        )

        takes_json_data = (
            endpoint.auth_header_content_type == 'application/json'
        )

        if date_is_skewed and takes_json_data:
            # On the real implementation, we get `codes.FORBIDDEN` and
            # `REQUEST_TIME_TOO_SKEWED`.
            # See https://github.com/adamtheturtle/vws-python/issues/407 for
            # implementing this on them mock.
            return

        if not date_is_skewed and takes_json_data:
            assert_vws_failure(
                response=response,
                status_code=codes.BAD_REQUEST,
                result_code=ResultCodes.FAIL,
            )
            return

        if date_is_skewed and is_summary_endpoint:
            assert_vws_failure(
                response=response,
                status_code=codes.FORBIDDEN,
                result_code=ResultCodes.REQUEST_TIME_TOO_SKEWED,
            )
            return

        if not date_is_skewed and is_summary_endpoint:
            assert_vws_failure(
                response=response,
                status_code=codes.UNAUTHORIZED,
                result_code=ResultCodes.AUTHENTICATION_FAILURE,
            )
            return

        assert response.status_code == codes.BAD_REQUEST
        url = str(endpoint.prepared_request.url)
        netloc = urlparse(url).netloc
        if netloc == 'cloudreco.vuforia.com':
            assert_vwq_failure(
                response=response,
                status_code=codes.BAD_REQUEST,
                content_type='text/html',
            )
            expected_text = (
                'java.lang.RuntimeException: RESTEASY007500: '
                'Could find no Content-Disposition header within part'
            )
            assert response.text == expected_text
            return

        assert response.text == ''
        assert 'Content-Type' not in response.headers
<|endoftext|>


<commit_before># Copyright (c) 2016 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

import ast
import math # Imported here so it can be used easily by the setting functions.

from UM.Logger import Logger

class IllegalMethodError(Exception):
    pass

##  Encapsulates Python code that provides a simple value calculation function.
#
class SettingFunction:
    ##  Constructor.
    #
    #   \param code The Python code this function should evaluate.
    def __init__(self, code, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._code = code
        self._settings = set() # Keys of all settings that are referenced to in this function.
        self._compiled = None
        self._valid = False

        try:
            tree = ast.parse(self._code, "eval")
            self._settings = set(_SettingExpressionVisitor().visit(tree))
            self._compiled = compile(self._code, repr(self), "eval")
            self._valid = True
        except (SyntaxError, TypeError) as e:
            Logger.log("e", "Parse error in function ({1}) for setting: {0}".format(str(e), self._code))
        except IllegalMethodError as e:
            Logger.log("e", "Use of illegal method {0} in function ({1}) for setting".format(str(e), self._code))
        except Exception as e:
            Logger.log("e", "Exception in function ({0}) for setting: {1}".format(str(e), self._code))

    ##  Call the actual function to calculate the value.
    def __call__(self, value_provider, *args, **kwargs):
        if not value_provider:
            return None

        if not self._valid:
            return None

        locals = { }
        for name in self._settings:
            locals[name] = value_provider.getProperty(name, "value")

        try:
            return eval(self._compiled, globals(), locals)
        except Exception as e:
            Logger.logException("d", "An exception occurred in inherit function %s", self)

    def __eq__(self, other):
        if not isinstance(other, SettingFunction):
            return False

        return self._code == other._code

    ##  Returns whether the function is ready to be executed.
    #
    #   \return True if the function is valid, or False if it's not.
    def isValid(self):
        return self._valid

    ##  Retrieve a set of the keys of all the settings used in this function.
    #
    #   \return A set of the keys of all the settings used in this functions.
    def getUsedSettings(self):
        return self._settings

    def __str__(self):
        return "SettingFunction({0})".format(self._code)

# Helper class used to analyze a parsed function
class _SettingExpressionVisitor(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.names = []

    def visit(self, node):
        super().visit(node)
        return self.names

    def visit_Name(self, node): # [CodeStyle: ast.NodeVisitor requires this function name]
        if node.id in self._blacklist:
            raise IllegalMethodError(node.id)

        if node.id not in self._knownNames and node.id not in __builtins__:
            self.names.append(node.id)

    _knownNames = [
        "math"
    ]

    _blacklist = [
        "sys",
        "os",
        "import",
        "__import__",
        "eval",
        "exec",
        "subprocess",
    ]
<commit_msg>Check for non-existant values when calling SettingFunction<commit_after># Copyright (c) 2016 Ultimaker B.V.
# Uranium is released under the terms of the AGPLv3 or higher.

import ast
import math # Imported here so it can be used easily by the setting functions.

from UM.Logger import Logger

class IllegalMethodError(Exception):
    pass

##  Encapsulates Python code that provides a simple value calculation function.
#
class SettingFunction:
    ##  Constructor.
    #
    #   \param code The Python code this function should evaluate.
    def __init__(self, code, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self._code = code
        self._settings = set() # Keys of all settings that are referenced to in this function.
        self._compiled = None
        self._valid = False

        try:
            tree = ast.parse(self._code, "eval")
            self._settings = set(_SettingExpressionVisitor().visit(tree))
            self._compiled = compile(self._code, repr(self), "eval")
            self._valid = True
        except (SyntaxError, TypeError) as e:
            Logger.log("e", "Parse error in function ({1}) for setting: {0}".format(str(e), self._code))
        except IllegalMethodError as e:
            Logger.log("e", "Use of illegal method {0} in function ({1}) for setting".format(str(e), self._code))
        except Exception as e:
            Logger.log("e", "Exception in function ({0}) for setting: {1}".format(str(e), self._code))

    ##  Call the actual function to calculate the value.
    def __call__(self, value_provider, *args, **kwargs):
        if not value_provider:
            return None

        if not self._valid:
            return None

        locals = { }
        for name in self._settings:
            value = value_provider.getProperty(name, "value")
            if value is None:
                Logger.log("e", "%s references unknown setting %s", self, name)
                continue

            locals[name] = value

        try:
            return eval(self._compiled, globals(), locals)
        except Exception as e:
            Logger.logException("d", "An exception occurred in inherit function %s", self)

    def __eq__(self, other):
        if not isinstance(other, SettingFunction):
            return False

        return self._code == other._code

    ##  Returns whether the function is ready to be executed.
    #
    #   \return True if the function is valid, or False if it's not.
    def isValid(self):
        return self._valid

    ##  Retrieve a set of the keys of all the settings used in this function.
    #
    #   \return A set of the keys of all the settings used in this functions.
    def getUsedSettings(self):
        return self._settings

    def __str__(self):
        return "SettingFunction({0})".format(self._code)

# Helper class used to analyze a parsed function
class _SettingExpressionVisitor(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.names = []

    def visit(self, node):
        super().visit(node)
        return self.names

    def visit_Name(self, node): # [CodeStyle: ast.NodeVisitor requires this function name]
        if node.id in self._blacklist:
            raise IllegalMethodError(node.id)

        if node.id not in self._knownNames and node.id not in __builtins__:
            self.names.append(node.id)

    _knownNames = [
        "math"
    ]

    _blacklist = [
        "sys",
        "os",
        "import",
        "__import__",
        "eval",
        "exec",
        "subprocess",
    ]
<|endoftext|>


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

def event_rank(count):
    """Determine event ranking."""
    if count < 10:
        return 'constable'
    elif count < 20:
        return 'sergeant'
    elif count < 50:
        return 'inspector'
    elif count < 100:
        return 'superintendent'
    elif count < 150:
        return 'commander'
    elif count > 200:
        return 'commissioner'

def users(store):
    """Display a list of all users and which institutes they belong to."""
    user_objs = list(store.users())
    total_events = store.user_events().count()
    for user_obj in user_objs:
        user_obj['institutes'] = [store.institute(inst_id) for inst_id in user_obj['institutes']]
        user_obj['events'] = store.user_events(user_obj).count()
        user_obj['events_rank'] = event_rank(user_obj['events'])
    return dict(
        users=user_objs,
        total_events=total_events,
    )
<commit_msg>adjust rankings for users<commit_after># -*- coding: utf-8 -*-

def event_rank(count):
    """Determine event ranking."""
    if count < 10:
        return 'constable'
    elif count < 100:
        return 'sergeant'
    elif count < 250:
        return 'inspector'
    elif count < 500:
        return 'superintendent'
    elif count < 1000:
        return 'commander'
    else:
        return 'commissioner'

def users(store):
    """Display a list of all users and which institutes they belong to."""
    user_objs = list(store.users())
    total_events = store.user_events().count()
    for user_obj in user_objs:
        user_obj['institutes'] = [store.institute(inst_id) for inst_id in user_obj['institutes']]
        user_obj['events'] = store.user_events(user_obj).count()
        user_obj['events_rank'] = event_rank(user_obj['events'])
    return dict(
        users=user_objs,
        total_events=total_events,
    )
<|endoftext|>


<commit_before>
"""Tools for Python database scripts."""

_symbols = {
    # skytools.timeutil
    'FixedOffsetTimezone': 'skytools.timeutil:FixedOffsetTimezone',
    'parse_iso_timestamp': 'skytools.timeutil:parse_iso_timestamp',
    # skytools.gzlog
    'gzip_append': 'skytools.gzlog:gzip_append',
    # skytools.config
    'Config': 'skytools.config:Config',
    # skytools.quoting
    'quote_bytea_copy': 'skytools.quoting:quote_bytea_copy',
    'quote_bytea_literal': 'skytools.quoting:quote_bytea_literal',
    'quote_bytea_raw': 'skytools.quoting:quote_bytea_raw',
    'quote_copy': 'skytools.quoting:quote_copy',
    'quote_fqident': 'skytools.quoting:quote_fqident',
    'quote_ident': 'skytools.quoting:quote_ident',
    'quote_json': 'skytools.quoting:quote_json',
    'quote_literal': 'skytools.quoting:quote_literal',
    'quote_statement': 'skytools.quoting:quote_statement',
    'db_urldecode': 'skytools.quoting:db_urldecode',
    'db_urlencode': 'skytools.quoting:db_urlencode',
    'unescape': 'skytools.quoting:unescape',
    'unescape_copy': 'skytools.quoting:unescape_copy',
    'unquote_fqident': 'skytools.quoting:unquote_fqident',
    'unquote_ident': 'skytools.quoting:unquote_ident',
    'unquote_literal': 'skytools.quoting:unquote_literal',
    'json_decode': 'skytools.quoting:json_decode',
    'json_encode': 'skytools.quoting:json_encode',
    # skytools.sqltools
    'dbdict': 'skytools.sqltools:dbdict',
    'CopyPipe': 'skytools.sqltools:CopyPipe',
    'DBObject': 'skytools.sqltools:DBObject',
    'DBFunction': 'skytools.sqltools:DBFunction',
    'DBLanguage': 'skytools.sqltools:DBLanguage',
    'DBSchema': 'skytools.sqltools:DBSchema',
    'DBTable': 'skytools.sqltools:DBTable',
    'Snapshot': 'skytools.sqltools:Snapshot',
    'db_install': 'skytools.sqltools:db_install',
    'exists_function': 'skytools.sqltools:exists_function',
    'exists_language': 'skytools.sqltools:exists_language',
    'exists_schema': 'skytools.sqltools:exists_schema',
    'exists_sequence': 'skytools.sqltools:exists_sequence',
    'exists_table': 'skytools.sqltools:exists_table',
    'exists_temp_table': 'skytools.sqltools:exists_temp_table',
    'exists_type': 'skytools.sqltools:exists_type',
    'fq_name': 'skytools.sqltools:fq_name',
    'fq_name_parts': 'skytools.sqltools:fq_name_parts',
    'full_copy': 'skytools.sqltools:full_copy',
    'get_table_columns': 'skytools.sqltools:get_table_columns',
    'get_table_oid': 'skytools.sqltools:get_table_oid',
    'get_table_pkeys': 'skytools.sqltools:get_table_pkeys',
    'installer_apply_file': 'skytools.sqltools:installer_apply_file',
    'installer_find_file': 'skytools.sqltools:installer_find_file',
    'magic_insert': 'skytools.sqltools:magic_insert',
    'mk_delete_sql': 'skytools.sqltools:mk_delete_sql',
    'mk_insert_sql': 'skytools.sqltools:mk_insert_sql',
    'mk_update_sql': 'skytools.sqltools:mk_update_sql',
    # skytools.psycopgwrapper
    'connect_database': 'skytools.psycopgwrapper:connect_database',
    'DBError': 'skytools.psycopgwrapper:DBError',
    # skytools.sockutil
    'set_tcp_keepalive': 'skytools.sockutil:set_tcp_keepalive',
    'set_cloexec': 'skytools.sockutil:set_cloexec',
    'set_nonblocking': 'skytools.sockutil:set_nonblocking',
    # skytools.scripting
    'BaseScript': 'skytools.scripting:BaseScript',
    'daemonize': 'skytools.scripting:daemonize',
    'DBScript': 'skytools.scripting:DBScript',
    'UsageError': 'skytools.scripting:UsageError',
    'signal_pidfile': 'skytools.scripting:signal_pidfile',
    'I_AUTOCOMMIT': 'skytools.scripting:I_AUTOCOMMIT',
    'I_READ_COMMITTED': 'skytools.scripting:I_READ_COMMITTED',
    'I_SERIALIZABLE': 'skytools.scripting:I_SERIALIZABLE',
    # skytools.adminscript
    'AdminScript': 'skytools.adminscript:AdminScript',
    # skytools.parsing
    'parse_acl': 'skytools.parsing:parse_acl',
    'parse_logtriga_sql': 'skytools.parsing:parse_logtriga_sql',
    'parse_pgarray': 'skytools.parsing:parse_pgarray',
    'parse_sqltriga_sql': 'skytools.parsing:parse_sqltriga_sql',
    'parse_statements': 'skytools.parsing:parse_statements',
    'parse_tabbed_table': 'skytools.parsing:parse_tabbed_table',
    'sql_tokenizer': 'skytools.parsing:sql_tokenizer',
    'dedent': 'skytools.parsing:dedent',
    # skytools.dbstruct
    'SeqStruct': 'skytools.dbstruct:SeqStruct',
    'TableStruct': 'skytools.dbstruct:TableStruct',
    'T_ALL': 'skytools.dbstruct:T_ALL',
    'T_CONSTRAINT': 'skytools.dbstruct:T_CONSTRAINT',
    'T_DEFAULT': 'skytools.dbstruct:T_DEFAULT',
    'T_GRANT': 'skytools.dbstruct:T_GRANT',
    'T_INDEX': 'skytools.dbstruct:T_INDEX',
    'T_OWNER': 'skytools.dbstruct:T_OWNER',
    'T_PARENT': 'skytools.dbstruct:T_PARENT',
    'T_PKEY': 'skytools.dbstruct:T_PKEY',
    'T_RULE': 'skytools.dbstruct:T_RULE',
    'T_SEQUENCE': 'skytools.dbstruct:T_SEQUENCE',
    'T_TABLE': 'skytools.dbstruct:T_TABLE',
    'T_TRIGGER': 'skytools.dbstruct:T_TRIGGER',
    # skytools.querybuilder
    'PLPyQuery': 'skytools.querybuilder:PLPyQuery',
    'QueryBuilder': 'skytools.querybuilder:QueryBuilder',
    'PLPyQueryBuilder': 'skytools.querybuilder:PLPyQueryBuilder',
    'plpy_exec': 'skytools.querybuilder:plpy_exec',
    'run_exists': 'skytools.querybuilder:run_exists',
    'run_lookup': 'skytools.querybuilder:run_lookup',
    'run_query': 'skytools.querybuilder:run_query',
    'run_query_row': 'skytools.querybuilder:run_query_row',
    # skytools.utf8
    'safe_utf8_decode': 'skytools.utf8:safe_utf8_decode',
}

__all__ = _symbols.keys()
_symbols['__version__'] = 'skytools.installer_config:package_version'

if 1:
    import skytools.apipkg as _apipkg
    _apipkg.initpkg(__name__, _symbols)
else:
    from skytools.timeutil import *
    from skytools.gzlog import *
    from skytools.config import *
    from skytools.quoting import *
    from skytools.parsing import *
    from skytools.sqltools import *
    from skytools.psycopgwrapper import *
    from skytools.dbstruct import *
    from skytools.scripting import *
    from skytools.adminscript import *
    from skytools.querybuilder import *

<commit_msg>skytools: make apipkg visible outside<commit_after>
"""Tools for Python database scripts."""

_symbols = {
    # skytools.timeutil
    'FixedOffsetTimezone': 'skytools.timeutil:FixedOffsetTimezone',
    'parse_iso_timestamp': 'skytools.timeutil:parse_iso_timestamp',
    # skytools.gzlog
    'gzip_append': 'skytools.gzlog:gzip_append',
    # skytools.config
    'Config': 'skytools.config:Config',
    # skytools.quoting
    'quote_bytea_copy': 'skytools.quoting:quote_bytea_copy',
    'quote_bytea_literal': 'skytools.quoting:quote_bytea_literal',
    'quote_bytea_raw': 'skytools.quoting:quote_bytea_raw',
    'quote_copy': 'skytools.quoting:quote_copy',
    'quote_fqident': 'skytools.quoting:quote_fqident',
    'quote_ident': 'skytools.quoting:quote_ident',
    'quote_json': 'skytools.quoting:quote_json',
    'quote_literal': 'skytools.quoting:quote_literal',
    'quote_statement': 'skytools.quoting:quote_statement',
    'db_urldecode': 'skytools.quoting:db_urldecode',
    'db_urlencode': 'skytools.quoting:db_urlencode',
    'unescape': 'skytools.quoting:unescape',
    'unescape_copy': 'skytools.quoting:unescape_copy',
    'unquote_fqident': 'skytools.quoting:unquote_fqident',
    'unquote_ident': 'skytools.quoting:unquote_ident',
    'unquote_literal': 'skytools.quoting:unquote_literal',
    'json_decode': 'skytools.quoting:json_decode',
    'json_encode': 'skytools.quoting:json_encode',
    # skytools.sqltools
    'dbdict': 'skytools.sqltools:dbdict',
    'CopyPipe': 'skytools.sqltools:CopyPipe',
    'DBObject': 'skytools.sqltools:DBObject',
    'DBFunction': 'skytools.sqltools:DBFunction',
    'DBLanguage': 'skytools.sqltools:DBLanguage',
    'DBSchema': 'skytools.sqltools:DBSchema',
    'DBTable': 'skytools.sqltools:DBTable',
    'Snapshot': 'skytools.sqltools:Snapshot',
    'db_install': 'skytools.sqltools:db_install',
    'exists_function': 'skytools.sqltools:exists_function',
    'exists_language': 'skytools.sqltools:exists_language',
    'exists_schema': 'skytools.sqltools:exists_schema',
    'exists_sequence': 'skytools.sqltools:exists_sequence',
    'exists_table': 'skytools.sqltools:exists_table',
    'exists_temp_table': 'skytools.sqltools:exists_temp_table',
    'exists_type': 'skytools.sqltools:exists_type',
    'fq_name': 'skytools.sqltools:fq_name',
    'fq_name_parts': 'skytools.sqltools:fq_name_parts',
    'full_copy': 'skytools.sqltools:full_copy',
    'get_table_columns': 'skytools.sqltools:get_table_columns',
    'get_table_oid': 'skytools.sqltools:get_table_oid',
    'get_table_pkeys': 'skytools.sqltools:get_table_pkeys',
    'installer_apply_file': 'skytools.sqltools:installer_apply_file',
    'installer_find_file': 'skytools.sqltools:installer_find_file',
    'magic_insert': 'skytools.sqltools:magic_insert',
    'mk_delete_sql': 'skytools.sqltools:mk_delete_sql',
    'mk_insert_sql': 'skytools.sqltools:mk_insert_sql',
    'mk_update_sql': 'skytools.sqltools:mk_update_sql',
    # skytools.psycopgwrapper
    'connect_database': 'skytools.psycopgwrapper:connect_database',
    'DBError': 'skytools.psycopgwrapper:DBError',
    # skytools.sockutil
    'set_tcp_keepalive': 'skytools.sockutil:set_tcp_keepalive',
    'set_cloexec': 'skytools.sockutil:set_cloexec',
    'set_nonblocking': 'skytools.sockutil:set_nonblocking',
    # skytools.scripting
    'BaseScript': 'skytools.scripting:BaseScript',
    'daemonize': 'skytools.scripting:daemonize',
    'DBScript': 'skytools.scripting:DBScript',
    'UsageError': 'skytools.scripting:UsageError',
    'signal_pidfile': 'skytools.scripting:signal_pidfile',
    'I_AUTOCOMMIT': 'skytools.scripting:I_AUTOCOMMIT',
    'I_READ_COMMITTED': 'skytools.scripting:I_READ_COMMITTED',
    'I_SERIALIZABLE': 'skytools.scripting:I_SERIALIZABLE',
    # skytools.adminscript
    'AdminScript': 'skytools.adminscript:AdminScript',
    # skytools.parsing
    'parse_acl': 'skytools.parsing:parse_acl',
    'parse_logtriga_sql': 'skytools.parsing:parse_logtriga_sql',
    'parse_pgarray': 'skytools.parsing:parse_pgarray',
    'parse_sqltriga_sql': 'skytools.parsing:parse_sqltriga_sql',
    'parse_statements': 'skytools.parsing:parse_statements',
    'parse_tabbed_table': 'skytools.parsing:parse_tabbed_table',
    'sql_tokenizer': 'skytools.parsing:sql_tokenizer',
    'dedent': 'skytools.parsing:dedent',
    # skytools.dbstruct
    'SeqStruct': 'skytools.dbstruct:SeqStruct',
    'TableStruct': 'skytools.dbstruct:TableStruct',
    'T_ALL': 'skytools.dbstruct:T_ALL',
    'T_CONSTRAINT': 'skytools.dbstruct:T_CONSTRAINT',
    'T_DEFAULT': 'skytools.dbstruct:T_DEFAULT',
    'T_GRANT': 'skytools.dbstruct:T_GRANT',
    'T_INDEX': 'skytools.dbstruct:T_INDEX',
    'T_OWNER': 'skytools.dbstruct:T_OWNER',
    'T_PARENT': 'skytools.dbstruct:T_PARENT',
    'T_PKEY': 'skytools.dbstruct:T_PKEY',
    'T_RULE': 'skytools.dbstruct:T_RULE',
    'T_SEQUENCE': 'skytools.dbstruct:T_SEQUENCE',
    'T_TABLE': 'skytools.dbstruct:T_TABLE',
    'T_TRIGGER': 'skytools.dbstruct:T_TRIGGER',
    # skytools.querybuilder
    'PLPyQuery': 'skytools.querybuilder:PLPyQuery',
    'QueryBuilder': 'skytools.querybuilder:QueryBuilder',
    'PLPyQueryBuilder': 'skytools.querybuilder:PLPyQueryBuilder',
    'plpy_exec': 'skytools.querybuilder:plpy_exec',
    'run_exists': 'skytools.querybuilder:run_exists',
    'run_lookup': 'skytools.querybuilder:run_lookup',
    'run_query': 'skytools.querybuilder:run_query',
    'run_query_row': 'skytools.querybuilder:run_query_row',
    # skytools.utf8
    'safe_utf8_decode': 'skytools.utf8:safe_utf8_decode',
}

__all__ = _symbols.keys()
_symbols['__version__'] = 'skytools.installer_config:package_version'

if 1:
    import skytools.apipkg as _apipkg
    _apipkg.initpkg(__name__, _symbols, {'apipkg': _apipkg})
else:
    from skytools.timeutil import *
    from skytools.gzlog import *
    from skytools.config import *
    from skytools.quoting import *
    from skytools.parsing import *
    from skytools.sqltools import *
    from skytools.psycopgwrapper import *
    from skytools.dbstruct import *
    from skytools.scripting import *
    from skytools.adminscript import *
    from skytools.querybuilder import *

<|endoftext|>