File: __main__.py

package info (click to toggle)
python-gsd 2.4.0-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 904 kB
  • sloc: python: 2,742; ansic: 1,881; makefile: 157; cpp: 109
file content (164 lines) | stat: -rw-r--r-- 5,054 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
# Copyright (c) 2020 The Regents of the University of Michigan
# All rights reserved.
# This software is licensed under the BSD 3-Clause License.

"""The GSD command line interface.

To simplify ad hoc usage of :py:mod:`gsd`, this module provides a command line
interface for interacting with GSD files. The primary entry point is a single
command for starting a Python interpreter with a GSD file pre-loaded::

    $ gsd read trajectory.gsd

The following options are available for the ``read`` subcommand:

.. program:: read

.. option:: -s schema, --schema schema

    The schema of the GSD file. Supported values for ``schema`` are "hoomd" and
    "none".

.. option:: -m mode, --mode mode

    The mode in which to open the file. Valid modes are identical to those
    accepted by :func:`gsd.fl.open`.
"""

import sys
import argparse
import code

from .version import __version__
from .hoomd import open as hoomd_open
from . import fl


def _print_err(msg=None, *args):
    print(msg, *args, file=sys.stderr)


SHELL_BANNER = """Python {python_version}
gsd {gsd_version}

File: {fn}
{extras}
The GSD file handle is available via the "handle" variable.
For supported schema, you may access the trajectory using the "traj" variable.
Type "help(handle)" or "help(traj)" for more information.
The gsd and gsd.fl packages are always loaded.
Schema-specific modules (e.g. gsd.hoomd) are loaded if available."""


def main_read(args):
    """Main function to launch a Python interpreter with an open GSD file."""
    # Default to a new line for well-formatted printing.
    local_ns = {
        'gsd': sys.modules['gsd'],
        'gsd.hoomd': sys.modules['gsd.hoomd'],
        'gsd.fl': sys.modules['gsd.fl'],
    }
    attributes = {}

    if args.schema == 'hoomd':
        traj = hoomd_open(args.file, mode=args.mode)
        handle = traj.file
        local_ns.update({
            'handle': handle,
            'traj': traj,
        })
        attributes.update({"Number of frames": len(traj)})
    else:
        if args.mode not in ['rb', 'rb+', 'ab']:
            raise ValueError("Unsupported schema for creating a file.")
        handle = fl.open(args.file, args.mode)
        local_ns.update({
            'handle': handle,
        })

    extras = "\n".join(
        "{}: {}".format(key, val) for key, val in attributes.items())

    code.interact(local=local_ns,
                  banner=SHELL_BANNER.format(python_version=sys.version,
                                             gsd_version=__version__,
                                             fn=args.file,
                                             extras=extras + "\n"))


def main():
    """Entry point to the GSD command-line interface.

    This function handles parsing command-line arguments and launching the
    appropriate subcommand based on the first argument to ``gsd`` on the
    command line. At present the following commands are supported:

        * read
    """
    parser = argparse.ArgumentParser(
        description="The gsd package encodes canonical readers and writers "
        "for the gsd file format.")
    parser.add_argument('--version',
                        action='store_true',
                        help="Display the version number and exit.")
    parser.add_argument('--debug',
                        action='store_true',
                        help="Show traceback on error for debugging.")
    subparsers = parser.add_subparsers()

    parser_read = subparsers.add_parser('read')
    parser_read.add_argument('file',
                             type=str,
                             nargs='?',
                             help="GSD file to read.")
    parser_read.add_argument('-s',
                             '--schema',
                             type=str,
                             default='hoomd',
                             choices=['hoomd', 'none'],
                             help="The file schema.")
    parser_read.add_argument(
        '-m',
        '--mode',
        type=str,
        default='rb',
        choices=['rb', 'rb+', 'wb', 'wb+', 'xb', 'xb+', 'ab'],
        help="The file mode.")
    parser_read.set_defaults(func=main_read)

    # This is a hack, as argparse itself does not
    # allow to parse only --version without any
    # of the other required arguments.
    if '--version' in sys.argv:
        print('gsd', __version__)
        sys.exit(0)

    args = parser.parse_args()

    if not hasattr(args, 'func'):
        parser.print_usage()
        sys.exit(2)
    try:
        args.func(args)
    except KeyboardInterrupt:
        _print_err()
        _print_err("Interrupted.")
        if args.debug:
            raise
        sys.exit(1)
    except RuntimeWarning as warning:
        _print_err("Warning: {}".format(warning))
        if args.debug:
            raise
        sys.exit(1)
    except Exception as error:
        _print_err('Error: {}'.format(error))
        if args.debug:
            raise
        sys.exit(1)
    else:
        sys.exit(0)


if __name__ == '__main__':
    main()