<commit_before>#!/usr/bin/env python
"""
Repackage a USGS Collection-1 tar for faster read access.

They arrive as a *.tar.gz with inner uncompressed tiffs, which Josh's tests have found to be too slow to read.

We compress the inner tiffs and store them in an uncompressed tar. This allows random reads within the files.
We also append a checksum file at the end of the tar.
"""
import copy
import io
import json
import stat
import tarfile
import tempfile
from functools import partial
from pathlib import Path
from typing import List, Iterable, Tuple, Callable, IO

import click
import numpy
import rasterio
from click import secho

from eodatasets.verify import PackageChecksum

_PREDICTOR_TABLE = {
    'int8': 2,
    'uint8': 2,
    'int16': 2,
    'uint16': 2,
    'int32': 2,
    'uint32': 2,
    'int64': 2,
    'uint64': 2,
    'float32': 3,
    'float64': 3
}

# The info of a file, and a method to open the file for reading.
ReadableMember = Tuple[tarfile.TarInfo, Callable[[], IO]]

def _create_tarinfo(path: Path, name=None) -> tarfile.TarInfo:
    """
    Create a TarInfo ("tar member") based on the given filesystem path.

    (these contain the information of a file, such as permissions, when writing to a tar file)

    This code is based on TarFile.gettarinfo(), but doesn't need an existing tar file.
    """
    # We're avoiding convenience methods like `path.is_file()`, to minimise repeated `stat()` calls on lustre.
    s = path.stat()
    info = tarfile.TarInfo(name or path.name)

    if stat.S_ISREG(s.st_mode):
        info.size = s.st_size
        info.type = tarfile.REGTYPE
    elif stat.S_ISDIR(s.st_mode):
        info.type = tarfile.DIRTYPE
        info.size = 0
    else:
        raise NotImplementedError(
            f"Only regular files and directories are supported for extracted datasets. "
            f"({path.name} in {path.absolute().parent})"
        )

    info.mode = s.st_mode
    info.uid = s.st_uid
    info.gid = s.st_gid
    info.mtime = s.st_mtime

    if tarfile.pwd:
        try:
            info.uname = tarfile.pwd.getpwuid(info.uid)[0]
        except KeyError:
            pass
    if tarfile.grp:
        try:
            info.gname = tarfile.grp.getgrgid(info.gid)[0]
        except KeyError:
            pass
    return info

def _tar_members(in_tar: tarfile.TarFile) -> Iterable[ReadableMember]:
    """Get readable files (members) from a tar"""
    members: List[tarfile.TarInfo] = in_tar.getmembers()

    for member in members:
        # We return a lambda/callable so that the file isn't opened until it's needed.
        yield member, partial(in_tar.extractfile, member)

def _folder_members(path: Path, base_path: Path = None) -> Iterable[ReadableMember]:
    """
    Get readable files (presented as tar members) from a directory.
    """
    if not base_path:
        base_path = path

    # Note that the members in input USGS tars are sorted alphabetically.
    # We'll sort our own inputs to match.
    # (The primary practical benefit is predictable outputs in tests)

    for item in sorted(path.iterdir()):
        member = _create_tarinfo(
            item,
            name=str(item.relative_to(base_path)),
        )
        if member.type == tarfile.DIRTYPE:
            yield member, None
            yield from _folder_members(item, base_path=path)
        else:
            # We return a lambda/callable so that the file isn't opened until it's needed.
            yield member, partial(item.open, 'rb')

def _create_tar_with_files(
        input_path: Path,
        input_files: Iterable[ReadableMember],
        output_tar_path: Path,
        **compress_args,
):
    """
    Package and compress the given input files to a new tar path.

    The output tar path is written atomically, so on failure it will only exist if complete.
    """
    if output_tar_path.exists():
        _log_skip(input_path, output_tar_path)
        return

    out_dir: Path = output_tar_path.parent
    out_dir.mkdir(parents=True, exist_ok=True)

    verify = PackageChecksum()

    # Use a temporary file so that we can move to the output path atomically.
    with tempfile.TemporaryDirectory(prefix='.extract-', dir=str(out_dir)) as tmpdir:
        tmpdir = Path(tmpdir).absolute()
        tmp_out_tar = tmpdir.joinpath(output_tar_path.name)

        with tarfile.open(tmp_out_tar, 'w') as out_tar:
            members = list(input_files)

            # Add the MTL file to the beginning of the output tar, so it can be accessed faster.
            # This slows down this repackage a little, as we're seeking/decompressing the input stream an extra time.
            _reorder_tar_members(members, input_path.name)

            with click.progressbar(label=input_path.name,
                                   length=sum(member.size for member, _ in members)) as progress:
                file_number = 0
                for member, open_member in members:
                    new_member = copy.copy(member)
                    # We'll copy them all as 664: matching USGS tars.
                    new_member.mode = 0o664

                    file_number += 1
                    progress.label = f"{input_path.name} ({file_number:2d}/{len(members)})"

                    # Recompress any TIFs, copy other files verbatim.
                    if member.name.lower().endswith('.tif'):
                        with rasterio.MemoryFile(filename=member.name) as memory_file:
                            try:
                                _recompress_image(open_member(), memory_file, **compress_args)
                            except Exception:
                                secho(f"Error during {member.name}", bold=True)
                                raise
                            new_member.size = memory_file.getbuffer().nbytes
                            out_tar.addfile(new_member, memory_file)

                            # Image has been written. Seek to beginning to take a checksum.
                            memory_file.seek(0)
                            verify.add(memory_file, tmpdir / new_member.name)
                    elif member.size == 0:
                        # Typically a directory entry.
                        out_tar.addfile(new_member)
                    else:
                        # Copy unchanged into target (typically the text/metadata files).
                        file_contents = open_member().read()
                        out_tar.addfile(new_member, io.BytesIO(file_contents))
                        verify.add(io.BytesIO(file_contents), tmpdir / new_member.name)
                        del file_contents

                    progress.update(member.size)

            # Append sha1 checksum file
            checksum_path = tmpdir / 'package.sha1'
            verify.write(checksum_path)
            checksum_path.chmod(0o664)
            out_tar.add(checksum_path, checksum_path.name)

            # Match the lower r/w permission bits to the output folder.
            # (Temp directories default to 700 otherwise.)
            tmp_out_tar.chmod(out_dir.stat().st_mode & 0o777)
            # Our output tar is complete. Move it into place.
            tmp_out_tar.rename(output_tar_path)

            _log_completion(members, input_path, output_tar_path)

def _log_skip(input_path: Path, output_tar: Path, reason='exists'):
    secho(
        json.dumps(
            dict(
                name=str(input_path.name),
                status=f'skip.{reason}',
                out_path=str(output_tar.absolute()),
                in_path=str(input_path.absolute()),
            )
        )
    )

def _log_completion(input_files: List[ReadableMember], input_path: Path, output_tar: Path):
    users = {(member.uname, member.gname) for member, _ in input_files}
    secho(
        json.dumps(
            dict(
                name=str(input_path.name),
                status='complete',
                in_size=sum(m.size for m, _ in input_files),
                in_count=len(input_files),
                # The user/group give us a hint as to whether this was repackaged outside of USGS.
                in_users=list(users),
                out_size=output_tar.stat().st_size,
                out_path=str(output_tar.absolute()),
                in_path=str(input_path.absolute()),
            )
        )
    )

def _reorder_tar_members(members: List[ReadableMember], identifier: str):
    """
    Put the (tiny) MTL file at the beginning of the tar so that it's always quick to read.
    """
    # Find MTL
    for i, (member, _) in enumerate(members):
        if '_MTL' in member.path:
            mtl_index = i
            break
    else:
        formatted_members = '\n\t'.join(m.name for m, _ in members)
        raise ValueError(f"No MTL file found in package {identifier}. Have:\n\t{formatted_members}")

    # Move to front
    mtl_item = members.pop(mtl_index)
    members.insert(0, mtl_item)

def _recompress_image(
        input_image_fp: IO,
        output_fp: rasterio.MemoryFile,
        zlevel=9,
        block_size=(512, 512),
):
    """
    Read an image from given file pointer, and write as a compressed GeoTIFF.
    """
    # noinspection PyUnusedLocal
    with rasterio.open(input_image_fp) as ds:
        ds: rasterio.DatasetReader
        block_size_y, block_size_x = block_size

        if len(ds.indexes) != 1:
            raise ValueError(f"Expecting one-band-per-tif input (USGS packages). "
                             f"Input has multiple layers {repr(ds.indexes)}")
        array: numpy.ndarray = ds.read(1)

        profile = ds.profile
        profile.update(
            driver='GTiff',
            predictor=_PREDICTOR_TABLE[array.dtype.name],
            compress='deflate',
            zlevel=zlevel,
            blockxsize=block_size_x,
            blockysize=block_size_y,
            tiled=True,
        )

        with output_fp.open(**profile) as output_dataset:
            output_dataset.write(array, 1)
            # Copy gdal metadata
            output_dataset.update_tags(**ds.tags())
            output_dataset.update_tags(1, **ds.tags(1))

@click.command(help=__doc__)
@click.option("--output-base", type=click.Path(file_okay=False, writable=True),
              help="The base output directory.", required=True)
@click.option("--zlevel", type=click.IntRange(0, 9), default=5,
              help="Deflate compression level.")
@click.option("--block-size", type=int, default=512,
              help="Compression block size (both x and y)")
@click.argument("paths", nargs=-1, type=click.Path(exists=True, readable=True))
def main(paths: List[str], output_base: str, zlevel: int, block_size: int):
    base_output_path = Path(output_base)
    with rasterio.Env():
        for path in paths:
            path = Path(path)

            try:
                # Input is either a tar.gz file, or a directory containing an MTL (already extracted)
                if path.suffix.lower() == '.gz':
                    with tarfile.open(str(path), 'r') as in_tar:
                        _create_tar_with_files(
                            path,
                            _tar_members(in_tar),
                            _output_tar_path(base_output_path, path),
                            zlevel=zlevel,
                            block_size=(block_size, block_size),
                        )
                elif path.is_dir():
                    _create_tar_with_files(
                        path,
                        _folder_members(path),
                        _output_tar_path_from_directory(base_output_path, path),
                        zlevel=zlevel,
                        block_size=(block_size, block_size),
                    )
                else:
                    raise ValueError(f"Expected either tar.gz or a dataset folder. Got: {path}")
            except Exception:
                secho(f"Error during {path}", color='red', bold=True)
                raise

def _output_tar_path(base_output, input_path):
    out_path = _calculate_out_base_path(base_output, input_path)
    return out_path.with_name(out_path.stem)

def _output_tar_path_from_directory(base_output, input_path):
    mtl_files = list(input_path.glob('*_MTL.txt'))
    if not mtl_files:
        raise ValueError(f"Dataset has no mtl: {input_path}")
    if len(mtl_files) > 1:
        secho(f"WARNING: multiple MTL files in {input_path}", bold=True, color='red', err=True)
    mtl_file = mtl_files[0]
    dataset_name = mtl_file.name.replace('_MTL.txt', '')
    out_tar_path = _calculate_out_base_path(base_output, input_path) / f'{dataset_name}.tar'
    return out_tar_path

def _calculate_out_base_path(out_base: Path, path: Path) -> Path:
    if 'USGS' not in path.parts:
        raise ValueError(
            "Expected AODH input path structure, "
            "eg: /AODH/USGS/L1/Landsat/C1/092_091/LT50920911991126/LT05_L1GS_092091_19910506_20170126_01_T2.tar.gz"
        )
    # The directory structure after the "USGS" folder is recreated onto the output base folder.
    return out_base.joinpath(*path.parts[path.parts.index('USGS') + 1:-1], path.name)

if __name__ == '__main__':
    main()
<commit_msg>Consistent logging output<commit_after>#!/usr/bin/env python
"""
Repackage a USGS Collection-1 tar for faster read access.

They arrive as a *.tar.gz with inner uncompressed tiffs, which Josh's tests have found to be too slow to read.

We compress the inner tiffs and store them in an uncompressed tar. This allows random reads within the files.
We also append a checksum file at the end of the tar.
"""
import copy
import io
import json
import stat
import sys
import tarfile
import tempfile
import traceback
from functools import partial
from pathlib import Path
from typing import List, Iterable, Tuple, Callable, IO

import click
import numpy
import rasterio
from click import secho, echo

from eodatasets.verify import PackageChecksum

_PREDICTOR_TABLE = {
    'int8': 2,
    'uint8': 2,
    'int16': 2,
    'uint16': 2,
    'int32': 2,
    'uint32': 2,
    'int64': 2,
    'uint64': 2,
    'float32': 3,
    'float64': 3
}

# The info of a file, and a method to open the file for reading.
ReadableMember = Tuple[tarfile.TarInfo, Callable[[], IO]]

def _create_tarinfo(path: Path, name=None) -> tarfile.TarInfo:
    """
    Create a TarInfo ("tar member") based on the given filesystem path.

    (these contain the information of a file, such as permissions, when writing to a tar file)

    This code is based on TarFile.gettarinfo(), but doesn't need an existing tar file.
    """
    # We're avoiding convenience methods like `path.is_file()`, to minimise repeated `stat()` calls on lustre.
    s = path.stat()
    info = tarfile.TarInfo(name or path.name)

    if stat.S_ISREG(s.st_mode):
        info.size = s.st_size
        info.type = tarfile.REGTYPE
    elif stat.S_ISDIR(s.st_mode):
        info.type = tarfile.DIRTYPE
        info.size = 0
    else:
        raise NotImplementedError(
            f"Only regular files and directories are supported for extracted datasets. "
            f"({path.name} in {path.absolute().parent})"
        )

    info.mode = s.st_mode
    info.uid = s.st_uid
    info.gid = s.st_gid
    info.mtime = s.st_mtime

    if tarfile.pwd:
        try:
            info.uname = tarfile.pwd.getpwuid(info.uid)[0]
        except KeyError:
            pass
    if tarfile.grp:
        try:
            info.gname = tarfile.grp.getgrgid(info.gid)[0]
        except KeyError:
            pass
    return info

def _tar_members(in_tar: tarfile.TarFile) -> Iterable[ReadableMember]:
    """Get readable files (members) from a tar"""
    members: List[tarfile.TarInfo] = in_tar.getmembers()

    for member in members:
        # We return a lambda/callable so that the file isn't opened until it's needed.
        yield member, partial(in_tar.extractfile, member)

def _folder_members(path: Path, base_path: Path = None) -> Iterable[ReadableMember]:
    """
    Get readable files (presented as tar members) from a directory.
    """
    if not base_path:
        base_path = path

    # Note that the members in input USGS tars are sorted alphabetically.
    # We'll sort our own inputs to match.
    # (The primary practical benefit is predictable outputs in tests)

    for item in sorted(path.iterdir()):
        member = _create_tarinfo(
            item,
            name=str(item.relative_to(base_path)),
        )
        if member.type == tarfile.DIRTYPE:
            yield member, None
            yield from _folder_members(item, base_path=path)
        else:
            # We return a lambda/callable so that the file isn't opened until it's needed.
            yield member, partial(item.open, 'rb')

def repackage_tar(
        input_path: Path,
        input_files: Iterable[ReadableMember],
        output_tar_path: Path,
        **compress_args,
) -> bool:
    if output_tar_path.exists():
        _log_skip(input_path, output_tar_path, 'exists')
        return True

    try:
        members = list(input_files)

        # Add the MTL file to the beginning of the output tar, so it can be accessed faster.
        # This slows down this repackage a little, as we're seeking/decompressing the input stream an extra time.
        _reorder_tar_members(members, input_path.name)

        _create_tar_with_files(input_path, members, output_tar_path, **compress_args)

        _log_completion(members, input_path, output_tar_path)
    except Exception as e:
        _log_skip(
            input_path,
            output_tar_path,
            'error',
            extra=dict(traceback=_format_exception(e))
        )
        return False
    return True

def _create_tar_with_files(
        input_path: Path,
        members: List[ReadableMember],
        output_tar_path: Path,
        **compress_args,
) -> None:
    """
    Package and compress the given input files to a new tar path.

    The output tar path is written atomically, so on failure it will only exist if complete.
    """

    out_dir: Path = output_tar_path.parent
    out_dir.mkdir(parents=True, exist_ok=True)

    verify = PackageChecksum()

    # Use a temporary file so that we can move to the output path atomically.
    with tempfile.TemporaryDirectory(prefix='.extract-', dir=str(out_dir)) as tmpdir:
        tmpdir = Path(tmpdir).absolute()
        tmp_out_tar = tmpdir.joinpath(output_tar_path.name)

        with tarfile.open(tmp_out_tar, 'w') as out_tar:
            with click.progressbar(label=input_path.name,
                                   length=sum(member.size for member, _ in members)) as progress:
                file_number = 0
                for member, open_member in members:
                    new_member = copy.copy(member)
                    # We'll copy them all as 664: matching USGS tars.
                    new_member.mode = 0o664

                    file_number += 1
                    progress.label = f"{input_path.name} ({file_number:2d}/{len(members)})"

                    # Recompress any TIFs, copy other files verbatim.
                    if member.name.lower().endswith('.tif'):
                        with rasterio.MemoryFile(filename=member.name) as memory_file:
                            try:
                                _recompress_image(open_member(), memory_file, **compress_args)
                            except Exception:
                                secho(f"Error during {member.name}", bold=True)
                                raise
                            new_member.size = memory_file.getbuffer().nbytes
                            out_tar.addfile(new_member, memory_file)

                            # Image has been written. Seek to beginning to take a checksum.
                            memory_file.seek(0)
                            verify.add(memory_file, tmpdir / new_member.name)
                    elif member.size == 0:
                        # Typically a directory entry.
                        out_tar.addfile(new_member)
                    else:
                        # Copy unchanged into target (typically the text/metadata files).
                        file_contents = open_member().read()
                        out_tar.addfile(new_member, io.BytesIO(file_contents))
                        verify.add(io.BytesIO(file_contents), tmpdir / new_member.name)
                        del file_contents

                    progress.update(member.size)

            # Append sha1 checksum file
            checksum_path = tmpdir / 'package.sha1'
            verify.write(checksum_path)
            checksum_path.chmod(0o664)
            out_tar.add(checksum_path, checksum_path.name)

            # Match the lower r/w permission bits to the output folder.
            # (Temp directories default to 700 otherwise.)
            tmp_out_tar.chmod(out_dir.stat().st_mode & 0o777)
            # Our output tar is complete. Move it into place.
            tmp_out_tar.rename(output_tar_path)

def _log_skip(input_path: Path, output_tar: Path, reason: str, extra: dict = None):
    secho(
        json.dumps(
            dict(
                name=str(input_path.name),
                status=f'skip.{reason}',
                out_path=str(output_tar.absolute()),
                in_path=str(input_path.absolute()),
                **(extra or {})
            )
        )
    )

def _log_completion(input_files: List[ReadableMember], input_path: Path, output_tar: Path):
    users = {(member.uname, member.gname) for member, _ in input_files}
    secho(
        json.dumps(
            dict(
                name=str(input_path.name),
                status='complete',
                in_size=sum(m.size for m, _ in input_files),
                in_count=len(input_files),
                # The user/group give us a hint as to whether this was repackaged outside of USGS.
                in_users=list(users),
                out_size=output_tar.stat().st_size,
                out_path=str(output_tar.absolute()),
                in_path=str(input_path.absolute()),
            )
        )
    )

def _reorder_tar_members(members: List[ReadableMember], identifier: str):
    """
    Put the (tiny) MTL file at the beginning of the tar so that it's always quick to read.
    """
    # Find MTL
    for i, (member, _) in enumerate(members):
        if '_MTL' in member.path:
            mtl_index = i
            break
    else:
        formatted_members = '\n\t'.join(m.name for m, _ in members)
        raise ValueError(f"No MTL file found in package {identifier}. Have:\n\t{formatted_members}")

    # Move to front
    mtl_item = members.pop(mtl_index)
    members.insert(0, mtl_item)

def _recompress_image(
        input_image_fp: IO,
        output_fp: rasterio.MemoryFile,
        zlevel=9,
        block_size=(512, 512),
):
    """
    Read an image from given file pointer, and write as a compressed GeoTIFF.
    """
    # noinspection PyUnusedLocal
    with rasterio.open(input_image_fp) as ds:
        ds: rasterio.DatasetReader
        block_size_y, block_size_x = block_size

        if len(ds.indexes) != 1:
            raise ValueError(f"Expecting one-band-per-tif input (USGS packages). "
                             f"Input has multiple layers {repr(ds.indexes)}")
        array: numpy.ndarray = ds.read(1)

        profile = ds.profile
        profile.update(
            driver='GTiff',
            predictor=_PREDICTOR_TABLE[array.dtype.name],
            compress='deflate',
            zlevel=zlevel,
            blockxsize=block_size_x,
            blockysize=block_size_y,
            tiled=True,
        )

        with output_fp.open(**profile) as output_dataset:
            output_dataset.write(array, 1)
            # Copy gdal metadata
            output_dataset.update_tags(**ds.tags())
            output_dataset.update_tags(1, **ds.tags(1))

@click.command(help=__doc__)
@click.option("--output-base", type=click.Path(file_okay=False, writable=True),
              help="The base output directory.", required=True)
@click.option("--zlevel", type=click.IntRange(0, 9), default=5,
              help="Deflate compression level.")
@click.option("--block-size", type=int, default=512,
              help="Compression block size (both x and y)")
@click.argument("paths", nargs=-1, type=click.Path(exists=True, readable=True))
def main(paths: List[str], output_base: str, zlevel: int, block_size: int):
    base_output_path = Path(output_base)
    with rasterio.Env():

        total = failures = 0
        for path in paths:
            total += 1
            path = Path(path)

            # Input is either a tar.gz file, or a directory containing an MTL (already extracted)
            if path.suffix.lower() == '.gz':
                with tarfile.open(str(path), 'r') as in_tar:
                    success = repackage_tar(
                        path,
                        _tar_members(in_tar),
                        _output_tar_path(base_output_path, path),
                        zlevel=zlevel,
                        block_size=(block_size, block_size),
                    )

            elif path.is_dir():
                success = repackage_tar(
                    path,
                    _folder_members(path),
                    _output_tar_path_from_directory(base_output_path, path),
                    zlevel=zlevel,
                    block_size=(block_size, block_size),
                )
            else:
                raise ValueError(f"Expected either tar.gz or a dataset folder. Got: {path}")

            if not success:
                failures += 1

    echo(f"Completed {total-failures} datasets. {failures} failures.")
    sys.exit(failures)

def _format_exception(e: BaseException):
    """
    Shamelessly stolen from stdlib's logging module.
    """
    sio = io.StringIO()
    traceback.print_exception(e.__class__, e, e.__traceback__, None, sio)
    s = sio.getvalue()
    sio.close()
    if s[-1:] == "\n":
        s = s[:-1]
    return s

def _output_tar_path(base_output, input_path):
    out_path = _calculate_out_base_path(base_output, input_path)
    return out_path.with_name(out_path.stem)

def _output_tar_path_from_directory(base_output, input_path):
    mtl_files = list(input_path.glob('*_MTL.txt'))
    if not mtl_files:
        raise ValueError(f"Dataset has no mtl: {input_path}")
    if len(mtl_files) > 1:
        secho(f"WARNING: multiple MTL files in {input_path}", bold=True, color='red', err=True)
    mtl_file = mtl_files[0]
    dataset_name = mtl_file.name.replace('_MTL.txt', '')
    out_tar_path = _calculate_out_base_path(base_output, input_path) / f'{dataset_name}.tar'
    return out_tar_path

def _calculate_out_base_path(out_base: Path, path: Path) -> Path:
    if 'USGS' not in path.parts:
        raise ValueError(
            "Expected AODH input path structure, "
            "eg: /AODH/USGS/L1/Landsat/C1/092_091/LT50920911991126/LT05_L1GS_092091_19910506_20170126_01_T2.tar.gz"
        )
    # The directory structure after the "USGS" folder is recreated onto the output base folder.
    return out_base.joinpath(*path.parts[path.parts.index('USGS') + 1:-1], path.name)

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


<commit_before>#***************************************************************************
#*                                                                         *
#*   Copyright (c) 2013 - Joachim Zettler                                  *
#*   Copyright (c) 2013 - Juergen Riegel <FreeCAD@juergen-riegel.net>      *
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   This program is distributed in the hope that it will be useful,       *
#*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
#*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
#*   GNU Library General Public License for more details.                  *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with this program; if not, write to the Free Software   *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************

import FreeCAD
import os
from math import pow, sqrt

__title__ = "FreeCAD Calculix library"
__author__ = "Juergen Riegel "
__url__ = "http://www.freecadweb.org"

if open.__module__ == '__builtin__':
    pyopen = open  # because we'll redefine open below

# read a calculix result file and extract the nodes, displacement vectores and stress values.
def readResult(frd_input):
    input = pyopen(frd_input, "r")
    nodes = {}
    disp = {}
    stress = {}
    elements = {}

    disp_found = False
    nodes_found = False
    stress_found = False
    elements_found = False
    elem = -1
    elemType = 0

    while True:
        line = input.readline()
        if not line:
            break
        #Check if we found nodes section
        if line[4:6] == "2C":
            nodes_found = True
        #first lets extract the node and coordinate information from the results file
        if nodes_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            nodes_x = float(line[13:25])
            nodes_y = float(line[25:37])
            nodes_z = float(line[37:49])
            nodes[elem] = FreeCAD.Vector(nodes_x, nodes_y, nodes_z)
        #Check if we found nodes section
        if line[4:6] == "3C":
            elements_found = True
        #first lets extract element number
        if elements_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            elemType = int(line[14:18])
        #then the 10 id's for the Tet10 element
        if elements_found and (line[1:3] == "-2") and elemType == 6:
            node_id_2 = int(line[3:13])
            node_id_1 = int(line[13:23])
            node_id_3 = int(line[23:33])
            node_id_4 = int(line[33:43])
            node_id_5 = int(line[43:53])
            node_id_7 = int(line[53:63])
            node_id_6 = int(line[63:73])
            node_id_9 = int(line[73:83])
            node_id_8 = int(line[83:93])
            node_id_10 = int(line[93:103])
            elements[elem] = (node_id_1, node_id_2, node_id_3, node_id_4, node_id_5, node_id_6, node_id_7, node_id_8, node_id_9, node_id_10)
        #Check if we found displacement section
        if line[5:9] == "DISP":
            disp_found = True
        #we found a displacement line in the frd file
        if disp_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            disp_x = float(line[13:25])
            disp_y = float(line[25:37])
            disp_z = float(line[37:49])
            disp[elem] = FreeCAD.Vector(disp_x, disp_y, disp_z)
        if line[5:11] == "STRESS":
            stress_found = True
        #we found a displacement line in the frd file
        if stress_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            stress_1 = float(line[13:25])
            stress_2 = float(line[25:37])
            stress_3 = float(line[37:49])
            stress_4 = float(line[49:61])
            stress_5 = float(line[61:73])
            stress_6 = float(line[73:85])
            stress[elem] = (stress_1, stress_2, stress_3, stress_4, stress_5, stress_6)
        #Check for the end of a section
        if line[1:3] == "-3":
            #the section with the displacements and the nodes ended
            disp_found = False
            nodes_found = False
            stress_found = False
            elements_found = False

    input.close()
    FreeCAD.Console.PrintLog('Read CalculiX result: {} Nodes, {} Displacements and {} Stress values\n'.format(len(nodes), len(disp), len(stress)))

    return {'Nodes': nodes, 'Tet10Elem': elements, 'Displacement': disp, 'Stress': stress}

def importFrd(filename, Analysis=None):
    m = readResult(filename)
    MeshObject = None
    if(len(m) > 0):
        import Fem
        if Analysis is None:
            AnalysisName = os.path.splitext(os.path.basename(filename))[0]
            AnalysisObject = FreeCAD.ActiveDocument.addObject('Fem::FemAnalysis', 'Analysis')
            AnalysisObject.Label = AnalysisName
        else:
            AnalysisObject = Analysis

        if ('Tet10Elem' in m) and ('Nodes' in m) and (not Analysis):
            mesh = Fem.FemMesh()
            nds = m['Nodes']
            for i in nds:
                n = nds[i]
                mesh.addNode(n[0], n[1], n[2], i)
            elms = m['Tet10Elem']
            for i in elms:
                e = elms[i]
                mesh.addVolume([e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9]], i)
            if len(nds) > 0:
                MeshObject = FreeCAD.ActiveDocument.addObject('Fem::FemMeshObject', 'ResultMesh')
                MeshObject.FemMesh = mesh
                AnalysisObject.Member = AnalysisObject.Member + [MeshObject]

        if 'Displacement' in m:
            disp = m['Displacement']
            if len(disp) > 0:
                o = FreeCAD.ActiveDocument.addObject('Fem::FemResultVector', 'Displacement')
                o.Values = disp.values()
                o.DataType = 'Displacement'
                o.ElementNumbers = disp.keys()
                if(MeshObject):
                    o.Mesh = MeshObject
                AnalysisObject.Member = AnalysisObject.Member + [o]
        if 'Stress' in m:
            stress = m['Stress']
            if len(stress) > 0:
                o = FreeCAD.ActiveDocument.addObject('Fem::FemResultValue', 'MisesStress')
                mstress = []
                for i in stress.values():
                    # Von mises stress (http://en.wikipedia.org/wiki/Von_Mises_yield_criterion)
                    s11 = i[0]
                    s22 = i[1]
                    s33 = i[2]
                    s12 = i[3]
                    s23 = i[4]
                    s31 = i[5]
                    s11s22 = pow(s11 - s22, 2)
                    s22s33 = pow(s22 - s33, 2)
                    s33s11 = pow(s33 - s11, 2)
                    s12s23s31 = 6 * (pow(s12, 2) + pow(s23, 2) * pow(s31, 2))
                    mstress.append(sqrt(0.5 * (s11s22 + s22s33 + s33s11 + s12s23s31)))

                o.Values = mstress
                o.DataType = 'VonMisesStress'
                o.ElementNumbers = stress.keys()
                if(MeshObject):
                    o.Mesh = MeshObject
                AnalysisObject.Member = AnalysisObject.Member + [o]
        if(FreeCAD.GuiUp):
            import FemGui
            import FreeCADGui
            if FreeCADGui.activeWorkbench().name() != 'FemWorkbench':
                FreeCADGui.activateWorkbench("FemWorkbench")
            FemGui.setActiveAnalysis(AnalysisObject)

def insert(filename, docname):
    "called when freecad wants to import a file"
    try:
        doc = FreeCAD.getDocument(docname)
    except NameError:
        doc = FreeCAD.newDocument(docname)
    FreeCAD.ActiveDocument = doc

    importFrd(filename)

def open(filename):
    "called when freecad opens a file"
    docname = os.path.splitext(os.path.basename(filename))[0]
    insert(filename, docname)
<commit_msg>FEM: replace input with frd_file in ccxFrdReader.py<commit_after>#***************************************************************************
#*                                                                         *
#*   Copyright (c) 2013 - Joachim Zettler                                  *
#*   Copyright (c) 2013 - Juergen Riegel <FreeCAD@juergen-riegel.net>      *
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   This program is distributed in the hope that it will be useful,       *
#*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
#*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
#*   GNU Library General Public License for more details.                  *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with this program; if not, write to the Free Software   *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************

import FreeCAD
import os
from math import pow, sqrt

__title__ = "FreeCAD Calculix library"
__author__ = "Juergen Riegel "
__url__ = "http://www.freecadweb.org"

if open.__module__ == '__builtin__':
    pyopen = open  # because we'll redefine open below

# read a calculix result file and extract the nodes, displacement vectores and stress values.
def readResult(frd_input):
    frd_file = pyopen(frd_input, "r")
    nodes = {}
    disp = {}
    stress = {}
    elements = {}

    disp_found = False
    nodes_found = False
    stress_found = False
    elements_found = False
    elem = -1
    elemType = 0

    while True:
        line = frd_file.readline()
        if not line:
            break
        #Check if we found nodes section
        if line[4:6] == "2C":
            nodes_found = True
        #first lets extract the node and coordinate information from the results file
        if nodes_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            nodes_x = float(line[13:25])
            nodes_y = float(line[25:37])
            nodes_z = float(line[37:49])
            nodes[elem] = FreeCAD.Vector(nodes_x, nodes_y, nodes_z)
        #Check if we found nodes section
        if line[4:6] == "3C":
            elements_found = True
        #first lets extract element number
        if elements_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            elemType = int(line[14:18])
        #then the 10 id's for the Tet10 element
        if elements_found and (line[1:3] == "-2") and elemType == 6:
            node_id_2 = int(line[3:13])
            node_id_1 = int(line[13:23])
            node_id_3 = int(line[23:33])
            node_id_4 = int(line[33:43])
            node_id_5 = int(line[43:53])
            node_id_7 = int(line[53:63])
            node_id_6 = int(line[63:73])
            node_id_9 = int(line[73:83])
            node_id_8 = int(line[83:93])
            node_id_10 = int(line[93:103])
            elements[elem] = (node_id_1, node_id_2, node_id_3, node_id_4, node_id_5, node_id_6, node_id_7, node_id_8, node_id_9, node_id_10)
        #Check if we found displacement section
        if line[5:9] == "DISP":
            disp_found = True
        #we found a displacement line in the frd file
        if disp_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            disp_x = float(line[13:25])
            disp_y = float(line[25:37])
            disp_z = float(line[37:49])
            disp[elem] = FreeCAD.Vector(disp_x, disp_y, disp_z)
        if line[5:11] == "STRESS":
            stress_found = True
        #we found a displacement line in the frd file
        if stress_found and (line[1:3] == "-1"):
            elem = int(line[4:13])
            stress_1 = float(line[13:25])
            stress_2 = float(line[25:37])
            stress_3 = float(line[37:49])
            stress_4 = float(line[49:61])
            stress_5 = float(line[61:73])
            stress_6 = float(line[73:85])
            stress[elem] = (stress_1, stress_2, stress_3, stress_4, stress_5, stress_6)
        #Check for the end of a section
        if line[1:3] == "-3":
            #the section with the displacements and the nodes ended
            disp_found = False
            nodes_found = False
            stress_found = False
            elements_found = False

    frd_file.close()
    FreeCAD.Console.PrintLog('Read CalculiX result: {} Nodes, {} Displacements and {} Stress values\n'.format(len(nodes), len(disp), len(stress)))

    return {'Nodes': nodes, 'Tet10Elem': elements, 'Displacement': disp, 'Stress': stress}

def importFrd(filename, Analysis=None):
    m = readResult(filename)
    MeshObject = None
    if(len(m) > 0):
        import Fem
        if Analysis is None:
            AnalysisName = os.path.splitext(os.path.basename(filename))[0]
            AnalysisObject = FreeCAD.ActiveDocument.addObject('Fem::FemAnalysis', 'Analysis')
            AnalysisObject.Label = AnalysisName
        else:
            AnalysisObject = Analysis

        if ('Tet10Elem' in m) and ('Nodes' in m) and (not Analysis):
            mesh = Fem.FemMesh()
            nds = m['Nodes']
            for i in nds:
                n = nds[i]
                mesh.addNode(n[0], n[1], n[2], i)
            elms = m['Tet10Elem']
            for i in elms:
                e = elms[i]
                mesh.addVolume([e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7], e[8], e[9]], i)
            if len(nds) > 0:
                MeshObject = FreeCAD.ActiveDocument.addObject('Fem::FemMeshObject', 'ResultMesh')
                MeshObject.FemMesh = mesh
                AnalysisObject.Member = AnalysisObject.Member + [MeshObject]

        if 'Displacement' in m:
            disp = m['Displacement']
            if len(disp) > 0:
                o = FreeCAD.ActiveDocument.addObject('Fem::FemResultVector', 'Displacement')
                o.Values = disp.values()
                o.DataType = 'Displacement'
                o.ElementNumbers = disp.keys()
                if(MeshObject):
                    o.Mesh = MeshObject
                AnalysisObject.Member = AnalysisObject.Member + [o]
        if 'Stress' in m:
            stress = m['Stress']
            if len(stress) > 0:
                o = FreeCAD.ActiveDocument.addObject('Fem::FemResultValue', 'MisesStress')
                mstress = []
                for i in stress.values():
                    # Von mises stress (http://en.wikipedia.org/wiki/Von_Mises_yield_criterion)
                    s11 = i[0]
                    s22 = i[1]
                    s33 = i[2]
                    s12 = i[3]
                    s23 = i[4]
                    s31 = i[5]
                    s11s22 = pow(s11 - s22, 2)
                    s22s33 = pow(s22 - s33, 2)
                    s33s11 = pow(s33 - s11, 2)
                    s12s23s31 = 6 * (pow(s12, 2) + pow(s23, 2) * pow(s31, 2))
                    mstress.append(sqrt(0.5 * (s11s22 + s22s33 + s33s11 + s12s23s31)))

                o.Values = mstress
                o.DataType = 'VonMisesStress'
                o.ElementNumbers = stress.keys()
                if(MeshObject):
                    o.Mesh = MeshObject
                AnalysisObject.Member = AnalysisObject.Member + [o]
        if(FreeCAD.GuiUp):
            import FemGui
            import FreeCADGui
            if FreeCADGui.activeWorkbench().name() != 'FemWorkbench':
                FreeCADGui.activateWorkbench("FemWorkbench")
            FemGui.setActiveAnalysis(AnalysisObject)

def insert(filename, docname):
    "called when freecad wants to import a file"
    try:
        doc = FreeCAD.getDocument(docname)
    except NameError:
        doc = FreeCAD.newDocument(docname)
    FreeCAD.ActiveDocument = doc

    importFrd(filename)

def open(filename):
    "called when freecad opens a file"
    docname = os.path.splitext(os.path.basename(filename))[0]
    insert(filename, docname)
<|endoftext|>


<commit_before>from coldtype import *

# Basics

# Pass an 300px inset rect containing an oval, with frame? taking a
# color tuple. What is a DATPen? <-- fontTools RecordingPen, a vector
# representation in Coldtype. DATPen vs P()??
#######################################################################
# @renderable((600, 600))
# def test(r:Rect):
#     #return DATPen().oval(r.inset(300)).f((0.3, 0.2, 0.6))
#     # return DATPen().rect(r.inset(300)).f(hsl(random()))
#     return P().rect(r.inset(50)).f(hsl(0.9))

# What is P()? P() == "pen, DATPen"
# f.a.r = frame.animation.rect, black, inset 300pixels
# with sinusoidal easing in/out
#######################################################################
# @animation(timeline=10) # duration of 10 frames
# def easing_example(f):
#     square = P(f.a.r.inset(300)).f(0)
#     return square.rotate(f.e("seio", 1, rng=(-3, 2)))

# Multiple Renderables
# Built in fonts are ColdtypeObviously, RecursiveMono, MutatorSans
#######################################################################
# @renderable((800, 600))
# def sample_text(r):
#     return PS([                                 # list for renderables
#         P().oval(r.inset(20)).f(hsl(random())), # renderable #1

#         (StSt("DOPO POTLY",                   # renderable #2
#         Font.ColdtypeObviously(), 200,
#         wdth=0, tu=100, rotate=15)
#         .align(r)
#         .f(1))])

# Less Basic Text
#######################################################################
# coldtype_obv = Font.ColdtypeObviously()
# @renderable((800,800))
# def basic(r):
#     # return (StSt("POTLY", coldtype_obv, 150).align(r))  # lessbasic
#     pens = (StSt("POTLY", coldtype_obv, 150,              # print_tree
#         wdth=0.5, rotate=10, tu=160)
#         .align(r)
#         # .f(hsl(0.8, s=0.75)))                           # no gradient
#         .f(Gradient.Vertical(r, hsl(0.5, s=0.8), hsl(0.8, s=0.75))))

#     print(pens.tree())
#     pens[0].rotate(180)
#     pens[-1].rotate(180)
#     pens[3].rotate(180)
#     return pens

# Drop Shadows
#######################################################################
coldtype_obv = Font.ColdtypeObviously()
# @renderable((1000, 200))
# def simpledrop(r):
#     pens = (StSt("LYPT CO", coldtype_obv, 150,
#         wdth=0.5, rotate=11, tu=250)
#         .align(r)
#         .f(1))
#     return DATPens([
#         pens.copy().translate(10, -10).f(0) # shadow
#         ,pens.s(hsl(0.9)).sw(3)             # top version of text
#     ])

# @renderable((1000, 200))
# def ro(r):
#     return (StSt("LPTT COOO", coldtype_obv, 150,
#         width=0.5, rotate=10, tu=100, ro=1)
#         .align(r)
#         .f(1)
#         .pen()
#         .layer(
#             lambda p: p.castshadow(-45, 50).f(0),
#             lambda p: p.s(hsl(0.9)).sw(3))
#         .align(r, th=1, tv=1))

# @renderable((1000, 200))
# def stroke_shadow(r):
#     return (StSt("LYYYT CO", coldtype_obv, 160,
#     width=1, rotate=10, tu=100, ro=1)
#     .align(r)
#     .f(1)
#     .layer(
#         lambda ps: ps.pmap(lambda p: p
#         .outline(10)
#         .removeOverlap()
#         .castshadow(-24, 50)
#         .f(None)
#         .s(hsl(0.6, s=1, l=0.4))
#         .sw(4)),
#         lambda ps: ps.s(hsl(0.9)).sw(4))
#     .align(r, th=1, tv=1)
#     )

@renderable((1000, 200))
def stroke_shadow_cleanup(r):
    def shadow_and_cleanup(p):
        return (p
            .outline(10)
            .reverse()
            .removeOverlap()
            .castshadow(-5, 220)
            # .filter_contours( lambda j, c:
            #     c.bounds().w > 50)
            .f(None)
            .s(hsl(0.6, s=1, l=0.4))
            .sw(4))

    return (StSt("C", coldtype_obv, 200,
        wdth=0.5, rotate=10, tu=100, ro=1)
        .align(r)
        .f(1)
        .layer(
            lambda ps: ps.pmap(shadow_and_cleanup),
            lambda ps: ps.s(hsl(0.9)).sw(4))
        .align(r, th=1, tv=1)
        )<commit_msg>[START] font trails for RR<commit_after>from coldtype import *
from coldtype.text.reader import offset
coldtype_obv = Font.ColdtypeObviously()

# Basics

# Pass an 300px inset rect containing an oval, with frame? taking a
# color tuple. What is a DATPen? <-- fontTools RecordingPen, a vector
# representation in Coldtype. DATPen vs P()??
#######################################################################
# @renderable((600, 600))
# def test(r:Rect):
#     #return DATPen().oval(r.inset(300)).f((0.3, 0.2, 0.6))
#     # return DATPen().rect(r.inset(300)).f(hsl(random()))
#     return P().rect(r.inset(50)).f(hsl(0.9))

# What is P()? P() == "pen, DATPen"
# f.a.r = frame.animation.rect, black, inset 300pixels
# with sinusoidal easing in/out
#######################################################################
# @animation(timeline=10) # duration of 10 frames
# def easing_example(f):
#     square = P(f.a.r.inset(300)).f(0)
#     return square.rotate(f.e("seio", 1, rng=(-3, 2)))

# Multiple Renderables
# Built in fonts are ColdtypeObviously, RecursiveMono, MutatorSans
#######################################################################
# @renderable((800, 600))
# def sample_text(r):
#     return PS([                                 # list for renderables
#         P().oval(r.inset(20)).f(hsl(random())), # renderable #1

#         (StSt("DOPO POTLY",                   # renderable #2
#         Font.ColdtypeObviously(), 200,
#         wdth=0, tu=100, rotate=15)
#         .align(r)
#         .f(1))])

# Less Basic Text
#######################################################################
# @renderable((800,800))
# def basic(r):
#     # return (StSt("POTLY", coldtype_obv, 150).align(r))  # lessbasic
#     pens = (StSt("POTLY", coldtype_obv, 150,              # print_tree
#         wdth=0.5, rotate=10, tu=160)
#         .align(r)
#         # .f(hsl(0.8, s=0.75)))                           # no gradient
#         .f(Gradient.Vertical(r, hsl(0.5, s=0.8), hsl(0.8, s=0.75))))

#     print(pens.tree())
#     pens[0].rotate(180)
#     pens[-1].rotate(180)
#     pens[3].rotate(180)
#     return pens

# Drop Shadows
#######################################################################
# @renderable((1000, 200))
# def simpledrop(r):
#     pens = (StSt("LYPT CO", coldtype_obv, 150,
#         wdth=0.5, rotate=11, tu=250)
#         .align(r)
#         .f(1))
#     return DATPens([
#         pens.copy().translate(10, -10).f(0) # shadow
#         ,pens.s(hsl(0.9)).sw(3)             # top version of text
#     ])

# @renderable((1000, 200))
# def ro(r):
#     return (StSt("LPTT COOO", coldtype_obv, 150,
#         width=0.5, rotate=10, tu=100, ro=1)
#         .align(r)
#         .f(1)
#         .pen()
#         .layer(
#             lambda p: p.castshadow(-45, 50).f(0),
#             lambda p: p.s(hsl(0.9)).sw(3))
#         .align(r, th=1, tv=1))

# @renderable((1000, 200))
# def stroke_shadow(r):
#     return (StSt("LYYYT CO", coldtype_obv, 160,
#     width=1, rotate=10, tu=100, ro=1)
#     .align(r)
#     .f(1)
#     .layer(
#         lambda ps: ps.pmap(lambda p: p
#         .outline(10)
#         .removeOverlap()
#         .castshadow(-24, 50)
#         .f(None)
#         .s(hsl(0.6, s=1, l=0.4))
#         .sw(4)),
#         lambda ps: ps.s(hsl(0.9)).sw(4))
#     .align(r, th=1, tv=1)
#     )

# cleanup a little spec in the above drop shadow rendering
# it may depend on the (canvas/char) dimensions chosen too, je pense
# @renderable((1000, 200))
# def stroke_shadow_cleanup(r):
#     def shadow_and_cleanup(p):
#         return (p
#             .outline(10)
#             .reverse()
#             .removeOverlap()
#             .castshadow(-5, 220)
#             # .filter_contours( lambda j, c:
#             #     c.bounds().w > 50)
#             .f(None)
#             .s(hsl(0.6, s=1, l=0.4))
#             .sw(4))

#     return (StSt("C", coldtype_obv, 200,
#         wdth=0.5, rotate=10, tu=100, ro=1)
#         .align(r)
#         .f(1)
#         .layer(
#             lambda ps: ps.pmap(shadow_and_cleanup),
#             lambda ps: ps.s(hsl(0.9)).sw(4))
#         .align(r, th=1, tv=1)
#         )

# Multi-line, just add a \n
#######################################################################
# @renderable((1000, 550))
# def multiline(r):
#     return (StSt("COLDTYPE\nTYPECOLD", coldtype_obv, 300,
#         wdth=1, fit=500)
#         .align(r)
#         .f(0))

# Text on a path
# @renderable((1000, 1000))
# def on_a_path(r):
#     circle = P().oval(r.inset(250)).reverse()
#     return (StSt("COLDTYPE", coldtype_obv, 200, width=1)
#         .distribute_on_path(circle, offset=275)
#         .f(0)
#     )

# Finally RR font experiments
rr_ttf = Font("rr_fonts/origami-mommy.pixellated.ttf")

#######################################################################
@renderable((1000, 550))
def multiline(r):
    return (StSt("RUOHO\nRECORDS", rr_ttf, 100,
        wdth=1, fit=500)
        .align(r)
        .f(0))<|endoftext|>


<commit_before>"""
The help command. The basic idea is that help texts for commands
are best written by those that write the commands - the admins. So
command-help is all auto-loaded and searched from the current command
set. The normal, database-tied help system is used for collaborative
creation of other help topics such as RP help or game-world aides.
"""

from django.conf import settings
from collections import defaultdict
from evennia.utils.utils import fill, dedent
from evennia.commands.command import Command
from evennia.help.models import HelpEntry
from evennia.utils import create
from evennia.utils.utils import string_suggestions, class_from_module

COMMAND_DEFAULT_CLASS = class_from_module(settings.COMMAND_DEFAULT_CLASS)

# limit symbol import for API
__all__ = ("CmdHelp", "CmdSetHelp")
_DEFAULT_WIDTH = settings.CLIENT_DEFAULT_WIDTH
_SEP = "{C" + "-" * _DEFAULT_WIDTH + "{n"

def format_help_entry(title, help_text, aliases=None, suggested=None):
    """
    This visually formats the help entry.
    """
    string = _SEP + "\n"
    if title:
        string += "{CHelp for {w%s{n" % title
    if aliases:
        string += " {C(aliases: %s{C){n" % ("{C,{n ".join("{w%s{n" % ali for ali in aliases))
    if help_text:
        string += "\n%s" % dedent(help_text.rstrip())
    if suggested:
        string += "\n\n{CSuggested:{n "
        string += "%s" % fill("{C,{n ".join("{w%s{n" % sug for sug in suggested))
    string.strip()
    string += "\n" + _SEP
    return string

def format_help_list(hdict_cmds, hdict_db):
    """
    Output a category-ordered list. The input are the
    pre-loaded help files for commands and database-helpfiles
    resectively.
    """
    string = ""
    if hdict_cmds and any(hdict_cmds.values()):
        string += "\n" + _SEP + "\n   {CCommand help entries{n\n" + _SEP
        for category in sorted(hdict_cmds.keys()):
            string += "\n  {w%s{n:\n" % (str(category).title())
            string += "{G" + fill(", ".join(sorted(hdict_cmds[category]))) + "{n"
    if hdict_db and any(hdict_db.values()):
        string += "\n\n" + _SEP + "\n\r  {COther help entries{n\n" + _SEP
        for category in sorted(hdict_db.keys()):
            string += "\n\r  {w%s{n:\n" % (str(category).title())
            string += "{G" + fill(", ".join(sorted([str(topic) for topic in hdict_db[category]]))) + "{n"
    return string

class CmdHelp(Command):
    """
    view help or a list of topics

    Usage:
      help <topic or command>
      help list
      help all

    This will search for help on commands and other
    topics related to the game.
    """
    key = "help"
    locks = "cmd:all()"
    arg_regex = r"\s|$"

    # this is a special cmdhandler flag that makes the cmdhandler also pack
    # the current cmdset with the call to self.func().
    return_cmdset = True

    def parse(self):
        """
        input is a string containing the command or topic to match.
        """
        self.original_args = self.args.strip()
        self.args = self.args.strip().lower()

    def check_show_help(self, cmd, caller):
        """
        Helper method. If this return True, the given cmd
        auto-help will be viewable in the help listing.
        Override this to easily select what is shown to
        the player. Note that only commands available
        in the caller's merged cmdset are available.

        Args:
            cmd (Command): Command class from the merged cmdset
            caller (Character, Player or Session): The current caller
                executing the help command.

        """
        # return only those with auto_help set and passing the cmd: lock
        return cmd.auto_help and cmd.access(caller)

    def func(self):
        """
        Run the dynamic help entry creator.
        """
        query, cmdset = self.args, self.cmdset
        caller = self.caller

        suggestion_cutoff = 0.6
        suggestion_maxnum = 5

        if not query:
            query = "all"

        # removing doublets in cmdset, caused by cmdhandler
        # having to allow doublet commands to manage exits etc.
        cmdset.make_unique(caller)

        # retrieve all available commands and database topics
        all_cmds = [cmd for cmd in cmdset if self.check_show_help(cmd, caller)]
        all_topics = [topic for topic in HelpEntry.objects.all() if topic.access(caller, 'view', default=True)]
        all_categories = list(set([cmd.help_category.lower() for cmd in all_cmds] + [topic.help_category.lower() for topic in all_topics]))

        if query in ("list", "all"):
            # we want to list all available help entries, grouped by category
            hdict_cmd = defaultdict(list)
            hdict_topic = defaultdict(list)
            # create the dictionaries {category:[topic, topic ...]} required by format_help_list
            [hdict_cmd[cmd.help_category].append(cmd.key) for cmd in all_cmds]
            [hdict_topic[topic.help_category].append(topic.key) for topic in all_topics]
            # report back
            self.msg(format_help_list(hdict_cmd, hdict_topic))
            return

        # Try to access a particular command

        # build vocabulary of suggestions and rate them by string similarity.
        vocabulary = [cmd.key for cmd in all_cmds if cmd] + [topic.key for topic in all_topics] + all_categories
        [vocabulary.extend(cmd.aliases) for cmd in all_cmds]
        suggestions = [sugg for sugg in string_suggestions(query, set(vocabulary), cutoff=suggestion_cutoff, maxnum=suggestion_maxnum)
                       if sugg != query]
        if not suggestions:
            suggestions = [sugg for sugg in vocabulary if sugg != query and sugg.startswith(query)]

        # try an exact command auto-help match
        match = [cmd for cmd in all_cmds if cmd == query]
        if len(match) == 1:
            self.msg(format_help_entry(match[0].key,
                     match[0].__doc__,
                     aliases=match[0].aliases,
                     suggested=suggestions))
            return

        # try an exact database help entry match
        match = list(HelpEntry.objects.find_topicmatch(query, exact=True))
        if len(match) == 1:
            self.msg(format_help_entry(match[0].key,
                     match[0].entrytext,
                     aliases=match[0].aliases.all(),
                     suggested=suggestions))
            return

        # try to see if a category name was entered
        if query in all_categories:
            self.msg(format_help_list({query:[cmd.key for cmd in all_cmds if cmd.help_category==query]},
                                        {query:[topic.key for topic in all_topics if topic.help_category==query]}))
            return

        # no exact matches found. Just give suggestions.
        self.msg(format_help_entry("", "No help entry found for '%s'" % query, None, suggested=suggestions))

class CmdSetHelp(COMMAND_DEFAULT_CLASS):
    """
    edit the help database

    Usage:
      @help[/switches] <topic>[[;alias;alias][,category[,locks]] = <text>

    Switches:
      replace - overwrite existing help topic.
      append - add text to the end of existing topic with a newline between.
      extend - as append, but don't add a newline.
      delete - remove help topic.

    Examples:
      @sethelp throw = This throws something at ...
      @sethelp/append pickpocketing,Thievery = This steals ...
      @sethelp/replace pickpocketing, ,attr(is_thief) = This steals ...

    This command manipulates the help database. A help entry can be created,
    appended/merged to and deleted. If you don't assign a category, the
    "General" category will be used. If no lockstring is specified, default
    is to let everyone read the help file.

    """
    key = "@help"
    aliases = "@sethelp"
    locks = "cmd:perm(PlayerHelpers)"
    help_category = "Building"

    def func(self):
        "Implement the function"

        switches = self.switches
        lhslist = self.lhslist

        if not self.args:
            self.msg("Usage: @sethelp[/switches] <topic>[;alias;alias][,category[,locks,..] = <text>")
            return

        nlist = len(lhslist)
        topicstr = lhslist[0] if nlist > 0 else ""
        if not topicstr:
            self.msg("You have to define a topic!")
            return
        topicstrlist = topicstr.split(";")
        topicstr, aliases = topicstrlist[0], topicstrlist[1:] if len(topicstr) > 1 else []
        aliastxt = ("(aliases: %s)" % ", ".join(aliases)) if aliases else ""
        old_entry = None

        # check if we have an old entry with the same name
        try:
            for querystr in topicstrlist:
                old_entry = HelpEntry.objects.find_topicmatch(querystr) # also search by alias
                if old_entry:
                    old_entry = list(old_entry)[0]
                    break
            category = lhslist[1] if nlist > 1 else old_entry.help_category
            lockstring = ",".join(lhslist[2:]) if nlist > 2 else old_entry.locks.get()
        except Exception:
            old_entry = None
            category = lhslist[1] if nlist > 1 else "General"
            lockstring = ",".join(lhslist[2:]) if nlist > 2 else "view:all()"
        category = category.lower()

        if 'append' in switches or "merge" in switches or "extend" in switches:
            # merge/append operations
            if not old_entry:
                self.msg("Could not find topic '%s'. You must give an exact name." % topicstr)
                return
            if not self.rhs:
                self.msg("You must supply text to append/merge.")
                return
            if 'merge' in switches:
                old_entry.entrytext += " " + self.rhs
            else:
                old_entry.entrytext += "\n%s" % self.rhs
            old_entry.aliases.add(aliases)
            self.msg("Entry updated:\n%s%s" % (old_entry.entrytext, aliastxt))
            return
        if 'delete' in switches or 'del' in switches:
            # delete the help entry
            if not old_entry:
                self.msg("Could not find topic '%s'%s." % (topicstr, aliastxt))
                return
            old_entry.delete()
            self.msg("Deleted help entry '%s'%s." % (topicstr, aliastxt))
            return

        # at this point it means we want to add a new help entry.
        if not self.rhs:
            self.msg("You must supply a help text to add.")
            return
        if old_entry:
            if 'replace' in switches:
                # overwrite old entry
                old_entry.key = topicstr
                old_entry.entrytext = self.rhs
                old_entry.help_category = category
                old_entry.locks.clear()
                old_entry.locks.add(lockstring)
                old_entry.aliases.add(aliases)
                old_entry.save()
                self.msg("Overwrote the old topic '%s'%s." % (topicstr, aliastxt))
            else:
                self.msg("Topic '%s'%s already exists. Use /replace to overwrite "
                        "or /append or /merge to add text to it." % (topicstr, aliastxt))
        else:
            # no old entry. Create a new one.
            new_entry = create.create_help_entry(topicstr,
                                                 self.rhs, category=category,
                                                 locks=lockstring,aliases=aliases)
            if new_entry:
                self.msg("Topic '%s'%s was successfully created." % (topicstr, aliastxt))
            else:
                self.msg("Error when creating topic '%s'%s! Contact an admin." % (topicstr, aliastxt))
<commit_msg>Add a few methods to CmdHelp, managing pagination of help messages (EvMore)<commit_after>"""
The help command. The basic idea is that help texts for commands
are best written by those that write the commands - the admins. So
command-help is all auto-loaded and searched from the current command
set. The normal, database-tied help system is used for collaborative
creation of other help topics such as RP help or game-world aides.
"""

from django.conf import settings
from collections import defaultdict
from evennia.utils.utils import fill, dedent
from evennia.commands.command import Command
from evennia.help.models import HelpEntry
from evennia.utils import create, evmore
from evennia.utils.utils import string_suggestions, class_from_module

COMMAND_DEFAULT_CLASS = class_from_module(settings.COMMAND_DEFAULT_CLASS)

# limit symbol import for API
__all__ = ("CmdHelp", "CmdSetHelp")
_DEFAULT_WIDTH = settings.CLIENT_DEFAULT_WIDTH
_SEP = "{C" + "-" * _DEFAULT_WIDTH + "{n"

class CmdHelp(Command):
    """
    View help or a list of topics

    Usage:
      help <topic or command>
      help list
      help all

    This will search for help on commands and other
    topics related to the game.
    """
    key = "help"
    locks = "cmd:all()"
    arg_regex = r"\s|$"

    # this is a special cmdhandler flag that makes the cmdhandler also pack
    # the current cmdset with the call to self.func().
    return_cmdset = True

    # Help messages are wrapped in an EvMore call.  If you want to
    # avoid this, simply set the 'help_more' flag to False.
    help_more = True

    @staticmethod
    def format_help_entry(title, help_text, aliases=None, suggested=None):
        """
        This visually formats the help entry.
        This method can be overriden to customize the way a help
        entry is displayed.

        Args:
            title (str): the title of the help entry.
            help_text (str): the text of the help entry.
            aliases (list of str or None): the list of aliases.
            suggested (list of str or None): suggested reading.

        Returns the formatted string, ready to be sent.

        """
        string = _SEP + "\n"
        if title:
            string += "{CHelp for {w%s{n" % title
        if aliases:
            string += " {C(aliases: %s{C){n" % ("{C,{n ".join("{w%s{n" % ali for ali in aliases))
        if help_text:
            string += "\n%s" % dedent(help_text.rstrip())
        if suggested:
            string += "\n\n{CSuggested:{n "
            string += "%s" % fill("{C,{n ".join("{w%s{n" % sug for sug in suggested))
        string.strip()
        string += "\n" + _SEP
        return string

    @staticmethod
    def format_help_list(hdict_cmds, hdict_db):
        """
        Output a category-ordered list. The input are the
        pre-loaded help files for commands and database-helpfiles
        respectively.  You can override this method to return a
        custom display of the list of commands and topics.
        """
        string = ""
        if hdict_cmds and any(hdict_cmds.values()):
            string += "\n" + _SEP + "\n   {CCommand help entries{n\n" + _SEP
            for category in sorted(hdict_cmds.keys()):
                string += "\n  {w%s{n:\n" % (str(category).title())
                string += "{G" + fill(", ".join(sorted(hdict_cmds[category]))) + "{n"
        if hdict_db and any(hdict_db.values()):
            string += "\n\n" + _SEP + "\n\r  {COther help entries{n\n" + _SEP
            for category in sorted(hdict_db.keys()):
                string += "\n\r  {w%s{n:\n" % (str(category).title())
                string += "{G" + fill(", ".join(sorted([str(topic) for topic in hdict_db[category]]))) + "{n"
        return string

    def check_show_help(self, cmd, caller):
        """
        Helper method. If this return True, the given cmd
        auto-help will be viewable in the help listing.
        Override this to easily select what is shown to
        the player. Note that only commands available
        in the caller's merged cmdset are available.

        Args:
            cmd (Command): Command class from the merged cmdset
            caller (Character, Player or Session): The current caller
                executing the help command.

        """
        # return only those with auto_help set and passing the cmd: lock
        return cmd.auto_help and cmd.access(caller)

    def should_list_cmd(self, cmd, caller):
        """
        Should the specified command appear in the help table?

        This method only checks whether a specified command should
        appear in the table of topics/commands.  The command can be
        used by the caller (see the 'check_show_help' method) and
        the command will still be available, for instance, if a
        character type 'help name of the command'.  However, if
        you return False, the specified command will not appear in
        the table.  This is sometimes useful to "hide" commands in
        the table, but still access them through the help system.

        Args:
            cmd: the command to be tested.
            caller: the caller of the help system.

        Return:
            True: the command should appear in the table.
            False: the command shouldn't appear in the table.

        """
        return True

    def parse(self):
        """
        input is a string containing the command or topic to match.
        """
        self.original_args = self.args.strip()
        self.args = self.args.strip().lower()

    def func(self):
        """
        Run the dynamic help entry creator.
        """
        query, cmdset = self.args, self.cmdset
        caller = self.caller

        suggestion_cutoff = 0.6
        suggestion_maxnum = 5

        if not query:
            query = "all"

        # removing doublets in cmdset, caused by cmdhandler
        # having to allow doublet commands to manage exits etc.
        cmdset.make_unique(caller)

        # retrieve all available commands and database topics
        all_cmds = [cmd for cmd in cmdset if self.check_show_help(cmd, caller)]
        all_topics = [topic for topic in HelpEntry.objects.all() if topic.access(caller, 'view', default=True)]
        all_categories = list(set([cmd.help_category.lower() for cmd in all_cmds] + [topic.help_category.lower() for topic in all_topics]))

        if query in ("list", "all"):
            # we want to list all available help entries, grouped by category
            hdict_cmd = defaultdict(list)
            hdict_topic = defaultdict(list)
            # create the dictionaries {category:[topic, topic ...]} required by format_help_list
            # Filter commands that should be reached by the help
            # system, but not be displayed in the table.
            for cmd in all_cmds:
                if self.should_list_cmd(cmd, caller):
                    hdict_cmd[cmd.help_category].append(cmd.key)
            [hdict_topic[topic.help_category].append(topic.key) for topic in all_topics]
            # report back
            self.msg(self.format_help_list(hdict_cmd, hdict_topic))
            return

        # Try to access a particular command

        # build vocabulary of suggestions and rate them by string similarity.
        vocabulary = [cmd.key for cmd in all_cmds if cmd] + [topic.key for topic in all_topics] + all_categories
        [vocabulary.extend(cmd.aliases) for cmd in all_cmds]
        suggestions = [sugg for sugg in string_suggestions(query, set(vocabulary), cutoff=suggestion_cutoff, maxnum=suggestion_maxnum)
                       if sugg != query]
        if not suggestions:
            suggestions = [sugg for sugg in vocabulary if sugg != query and sugg.startswith(query)]

        # try an exact command auto-help match
        match = [cmd for cmd in all_cmds if cmd == query]
        if len(match) == 1:
            formatted = self.format_help_entry(match[0].key,
                     match[0].__doc__,
                     aliases=match[0].aliases,
                     suggested=suggestions)
            if type(self).help_more:
                evmore.msg(caller, formatted)
            else:
                self.msg(formatted)
            return

        # try an exact database help entry match
        match = list(HelpEntry.objects.find_topicmatch(query, exact=True))
        if len(match) == 1:
            formatted = self.format_help_entry(match[0].key,
                     match[0].entrytext,
                     aliases=match[0].aliases.all(),
                     suggested=suggestions)
            if type(self).help_more:
                evmore.msg(caller, formatted)
            else:
                self.msg(formatted)
            return

        # try to see if a category name was entered
        if query in all_categories:
            self.msg(self.format_help_list({query:[cmd.key for cmd in all_cmds if cmd.help_category==query]},
                                        {query:[topic.key for topic in all_topics if topic.help_category==query]}))
            return

        # no exact matches found. Just give suggestions.
        self.msg(format_help_entry("", "No help entry found for '%s'" % query, None, suggested=suggestions))

class CmdSetHelp(COMMAND_DEFAULT_CLASS):
    """
    edit the help database

    Usage:
      @help[/switches] <topic>[[;alias;alias][,category[,locks]] = <text>

    Switches:
      replace - overwrite existing help topic.
      append - add text to the end of existing topic with a newline between.
      extend - as append, but don't add a newline.
      delete - remove help topic.

    Examples:
      @sethelp throw = This throws something at ...
      @sethelp/append pickpocketing,Thievery = This steals ...
      @sethelp/replace pickpocketing, ,attr(is_thief) = This steals ...

    This command manipulates the help database. A help entry can be created,
    appended/merged to and deleted. If you don't assign a category, the
    "General" category will be used. If no lockstring is specified, default
    is to let everyone read the help file.

    """
    key = "@help"
    aliases = "@sethelp"
    locks = "cmd:perm(PlayerHelpers)"
    help_category = "Building"

    def func(self):
        "Implement the function"

        switches = self.switches
        lhslist = self.lhslist

        if not self.args:
            self.msg("Usage: @sethelp[/switches] <topic>[;alias;alias][,category[,locks,..] = <text>")
            return

        nlist = len(lhslist)
        topicstr = lhslist[0] if nlist > 0 else ""
        if not topicstr:
            self.msg("You have to define a topic!")
            return
        topicstrlist = topicstr.split(";")
        topicstr, aliases = topicstrlist[0], topicstrlist[1:] if len(topicstr) > 1 else []
        aliastxt = ("(aliases: %s)" % ", ".join(aliases)) if aliases else ""
        old_entry = None

        # check if we have an old entry with the same name
        try:
            for querystr in topicstrlist:
                old_entry = HelpEntry.objects.find_topicmatch(querystr) # also search by alias
                if old_entry:
                    old_entry = list(old_entry)[0]
                    break
            category = lhslist[1] if nlist > 1 else old_entry.help_category
            lockstring = ",".join(lhslist[2:]) if nlist > 2 else old_entry.locks.get()
        except Exception:
            old_entry = None
            category = lhslist[1] if nlist > 1 else "General"
            lockstring = ",".join(lhslist[2:]) if nlist > 2 else "view:all()"
        category = category.lower()

        if 'append' in switches or "merge" in switches or "extend" in switches:
            # merge/append operations
            if not old_entry:
                self.msg("Could not find topic '%s'. You must give an exact name." % topicstr)
                return
            if not self.rhs:
                self.msg("You must supply text to append/merge.")
                return
            if 'merge' in switches:
                old_entry.entrytext += " " + self.rhs
            else:
                old_entry.entrytext += "\n%s" % self.rhs
            old_entry.aliases.add(aliases)
            self.msg("Entry updated:\n%s%s" % (old_entry.entrytext, aliastxt))
            return
        if 'delete' in switches or 'del' in switches:
            # delete the help entry
            if not old_entry:
                self.msg("Could not find topic '%s'%s." % (topicstr, aliastxt))
                return
            old_entry.delete()
            self.msg("Deleted help entry '%s'%s." % (topicstr, aliastxt))
            return

        # at this point it means we want to add a new help entry.
        if not self.rhs:
            self.msg("You must supply a help text to add.")
            return
        if old_entry:
            if 'replace' in switches:
                # overwrite old entry
                old_entry.key = topicstr
                old_entry.entrytext = self.rhs
                old_entry.help_category = category
                old_entry.locks.clear()
                old_entry.locks.add(lockstring)
                old_entry.aliases.add(aliases)
                old_entry.save()
                self.msg("Overwrote the old topic '%s'%s." % (topicstr, aliastxt))
            else:
                self.msg("Topic '%s'%s already exists. Use /replace to overwrite "
                        "or /append or /merge to add text to it." % (topicstr, aliastxt))
        else:
            # no old entry. Create a new one.
            new_entry = create.create_help_entry(topicstr,
                                                 self.rhs, category=category,
                                                 locks=lockstring,aliases=aliases)
            if new_entry:
                self.msg("Topic '%s'%s was successfully created." % (topicstr, aliastxt))
            else:
                self.msg("Error when creating topic '%s'%s! Contact an admin." % (topicstr, aliastxt))
<|endoftext|>