hgclient.py
134 lines
| 3.6 KiB
| text/x-python
|
PythonLexer
/ contrib / hgclient.py
|
r22566 | # A minimal client for Mercurial's command server | ||
|
r28355 | from __future__ import absolute_import, print_function | ||
|
r40353 | |||
import io | ||||
|
r28355 | import os | ||
|
r40352 | import re | ||
|
r28355 | import signal | ||
import socket | ||||
import struct | ||||
import subprocess | ||||
import sys | ||||
import time | ||||
|
r22566 | |||
|
r40351 | if sys.version_info[0] >= 3: | ||
stdout = sys.stdout.buffer | ||||
stderr = sys.stderr.buffer | ||||
|
r40353 | stringio = io.BytesIO | ||
|
r40352 | def bprint(*args): | ||
# remove b'' as well for ease of test migration | ||||
pargs = [re.sub(br'''\bb(['"])''', br'\1', b'%s' % a) for a in args] | ||||
stdout.write(b' '.join(pargs) + b'\n') | ||||
|
r40351 | else: | ||
|
r40353 | import cStringIO | ||
|
r40351 | stdout = sys.stdout | ||
stderr = sys.stderr | ||||
|
r40353 | stringio = cStringIO.StringIO | ||
|
r40352 | bprint = print | ||
|
r40351 | |||
|
r22992 | def connectpipe(path=None): | ||
|
r40350 | cmdline = [b'hg', b'serve', b'--cmdserver', b'pipe'] | ||
|
r22566 | if path: | ||
|
r40350 | cmdline += [b'-R', path] | ||
|
r22566 | |||
server = subprocess.Popen(cmdline, stdin=subprocess.PIPE, | ||||
stdout=subprocess.PIPE) | ||||
return server | ||||
|
r22993 | class unixconnection(object): | ||
def __init__(self, sockpath): | ||||
self.sock = sock = socket.socket(socket.AF_UNIX) | ||||
sock.connect(sockpath) | ||||
self.stdin = sock.makefile('wb') | ||||
self.stdout = sock.makefile('rb') | ||||
def wait(self): | ||||
self.stdin.close() | ||||
self.stdout.close() | ||||
self.sock.close() | ||||
class unixserver(object): | ||||
def __init__(self, sockpath, logpath=None, repopath=None): | ||||
self.sockpath = sockpath | ||||
|
r40350 | cmdline = [b'hg', b'serve', b'--cmdserver', b'unix', b'-a', sockpath] | ||
|
r22993 | if repopath: | ||
|
r40350 | cmdline += [b'-R', repopath] | ||
|
r22993 | if logpath: | ||
stdout = open(logpath, 'a') | ||||
stderr = subprocess.STDOUT | ||||
else: | ||||
stdout = stderr = None | ||||
self.server = subprocess.Popen(cmdline, stdout=stdout, stderr=stderr) | ||||
# wait for listen() | ||||
while self.server.poll() is None: | ||||
if os.path.exists(sockpath): | ||||
break | ||||
time.sleep(0.1) | ||||
def connect(self): | ||||
return unixconnection(self.sockpath) | ||||
def shutdown(self): | ||||
os.kill(self.server.pid, signal.SIGTERM) | ||||
self.server.wait() | ||||
|
r22566 | def writeblock(server, data): | ||
|
r40350 | server.stdin.write(struct.pack(b'>I', len(data))) | ||
|
r22566 | server.stdin.write(data) | ||
server.stdin.flush() | ||||
def readchannel(server): | ||||
data = server.stdout.read(5) | ||||
if not data: | ||||
raise EOFError | ||||
channel, length = struct.unpack('>cI', data) | ||||
|
r40350 | if channel in b'IL': | ||
|
r22566 | return channel, length | ||
else: | ||||
return channel, server.stdout.read(length) | ||||
def sep(text): | ||||
|
r40350 | return text.replace(b'\\', b'/') | ||
|
r22566 | |||
|
r40351 | def runcommand(server, args, output=stdout, error=stderr, input=None, | ||
|
r22566 | outfilter=lambda x: x): | ||
|
r40352 | bprint(b'*** runcommand', b' '.join(args)) | ||
|
r40351 | stdout.flush() | ||
|
r40350 | server.stdin.write(b'runcommand\n') | ||
writeblock(server, b'\0'.join(args)) | ||||
|
r22566 | |||
if not input: | ||||
|
r28836 | input = stringio() | ||
|
r22566 | |||
while True: | ||||
ch, data = readchannel(server) | ||||
|
r40350 | if ch == b'o': | ||
|
r22566 | output.write(outfilter(data)) | ||
output.flush() | ||||
|
r40350 | elif ch == b'e': | ||
|
r22566 | error.write(data) | ||
error.flush() | ||||
|
r40350 | elif ch == b'I': | ||
|
r22566 | writeblock(server, input.read(data)) | ||
|
r40350 | elif ch == b'L': | ||
|
r22566 | writeblock(server, input.readline(data)) | ||
|
r40350 | elif ch == b'r': | ||
|
r22566 | ret, = struct.unpack('>i', data) | ||
if ret != 0: | ||||
|
r40352 | bprint(b' [%d]' % ret) | ||
|
r22566 | return ret | ||
else: | ||||
|
r40352 | bprint(b"unexpected channel %c: %r" % (ch, data)) | ||
|
r22566 | if ch.isupper(): | ||
return | ||||
|
r22992 | def check(func, connect=connectpipe): | ||
|
r40351 | stdout.flush() | ||
|
r22991 | server = connect() | ||
|
r22566 | try: | ||
return func(server) | ||||
finally: | ||||
server.stdin.close() | ||||
server.wait() | ||||