hgclient.py
161 lines
| 3.9 KiB
| text/x-python
|
PythonLexer
/ contrib / hgclient.py
Yuya Nishihara
|
r22566 | # A minimal client for Mercurial's command server | ||
Yuya Nishihara
|
r40353 | |||
import io | ||||
Pulkit Goyal
|
r28355 | import os | ||
Yuya Nishihara
|
r40352 | import re | ||
Pulkit Goyal
|
r28355 | import signal | ||
import socket | ||||
import struct | ||||
import subprocess | ||||
import sys | ||||
import time | ||||
Yuya Nishihara
|
r22566 | |||
Yuya Nishihara
|
r40351 | if sys.version_info[0] >= 3: | ||
stdout = sys.stdout.buffer | ||||
stderr = sys.stderr.buffer | ||||
Yuya Nishihara
|
r40353 | stringio = io.BytesIO | ||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
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') | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r40351 | else: | ||
Yuya Nishihara
|
r40353 | import cStringIO | ||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r40351 | stdout = sys.stdout | ||
stderr = sys.stderr | ||||
Yuya Nishihara
|
r40353 | stringio = cStringIO.StringIO | ||
Yuya Nishihara
|
r40352 | bprint = print | ||
Yuya Nishihara
|
r40351 | |||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r40625 | def connectpipe(path=None, extraargs=()): | ||
Yuya Nishihara
|
r40350 | cmdline = [b'hg', b'serve', b'--cmdserver', b'pipe'] | ||
Yuya Nishihara
|
r22566 | if path: | ||
Yuya Nishihara
|
r40350 | cmdline += [b'-R', path] | ||
Yuya Nishihara
|
r40625 | cmdline.extend(extraargs) | ||
Yuya Nishihara
|
r22566 | |||
Matt Harbison
|
r41020 | def tonative(cmdline): | ||
Augie Fackler
|
r43906 | if os.name != 'nt': | ||
Matt Harbison
|
r41020 | return cmdline | ||
return [arg.decode("utf-8") for arg in cmdline] | ||||
Augie Fackler
|
r43346 | server = subprocess.Popen( | ||
tonative(cmdline), stdin=subprocess.PIPE, stdout=subprocess.PIPE | ||||
) | ||||
Yuya Nishihara
|
r22566 | |||
return server | ||||
Augie Fackler
|
r43346 | |||
Gregory Szorc
|
r49801 | class unixconnection: | ||
Yuya Nishihara
|
r22993 | 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() | ||||
Augie Fackler
|
r43346 | |||
Gregory Szorc
|
r49801 | class unixserver: | ||
Yuya Nishihara
|
r22993 | def __init__(self, sockpath, logpath=None, repopath=None): | ||
self.sockpath = sockpath | ||||
Yuya Nishihara
|
r40350 | cmdline = [b'hg', b'serve', b'--cmdserver', b'unix', b'-a', sockpath] | ||
Yuya Nishihara
|
r22993 | if repopath: | ||
Yuya Nishihara
|
r40350 | cmdline += [b'-R', repopath] | ||
Yuya Nishihara
|
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() | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r22566 | def writeblock(server, data): | ||
Yuya Nishihara
|
r40350 | server.stdin.write(struct.pack(b'>I', len(data))) | ||
Yuya Nishihara
|
r22566 | server.stdin.write(data) | ||
server.stdin.flush() | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r22566 | def readchannel(server): | ||
data = server.stdout.read(5) | ||||
if not data: | ||||
raise EOFError | ||||
channel, length = struct.unpack('>cI', data) | ||||
Yuya Nishihara
|
r40350 | if channel in b'IL': | ||
Yuya Nishihara
|
r22566 | return channel, length | ||
else: | ||||
return channel, server.stdout.read(length) | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r22566 | def sep(text): | ||
Yuya Nishihara
|
r40350 | return text.replace(b'\\', b'/') | ||
Yuya Nishihara
|
r22566 | |||
Augie Fackler
|
r43346 | |||
def runcommand( | ||||
server, args, output=stdout, error=stderr, input=None, outfilter=lambda x: x | ||||
): | ||||
Yuya Nishihara
|
r40352 | bprint(b'*** runcommand', b' '.join(args)) | ||
Yuya Nishihara
|
r40351 | stdout.flush() | ||
Yuya Nishihara
|
r40350 | server.stdin.write(b'runcommand\n') | ||
writeblock(server, b'\0'.join(args)) | ||||
Yuya Nishihara
|
r22566 | |||
if not input: | ||||
timeless
|
r28836 | input = stringio() | ||
Yuya Nishihara
|
r22566 | |||
while True: | ||||
ch, data = readchannel(server) | ||||
Yuya Nishihara
|
r40350 | if ch == b'o': | ||
Yuya Nishihara
|
r22566 | output.write(outfilter(data)) | ||
output.flush() | ||||
Yuya Nishihara
|
r40350 | elif ch == b'e': | ||
Yuya Nishihara
|
r22566 | error.write(data) | ||
error.flush() | ||||
Yuya Nishihara
|
r40350 | elif ch == b'I': | ||
Yuya Nishihara
|
r22566 | writeblock(server, input.read(data)) | ||
Yuya Nishihara
|
r40350 | elif ch == b'L': | ||
Yuya Nishihara
|
r22566 | writeblock(server, input.readline(data)) | ||
Yuya Nishihara
|
r40625 | elif ch == b'm': | ||
bprint(b"message: %r" % data) | ||||
Yuya Nishihara
|
r40350 | elif ch == b'r': | ||
Augie Fackler
|
r43346 | (ret,) = struct.unpack('>i', data) | ||
Yuya Nishihara
|
r22566 | if ret != 0: | ||
Yuya Nishihara
|
r40352 | bprint(b' [%d]' % ret) | ||
Yuya Nishihara
|
r22566 | return ret | ||
else: | ||||
Yuya Nishihara
|
r40352 | bprint(b"unexpected channel %c: %r" % (ch, data)) | ||
Yuya Nishihara
|
r22566 | if ch.isupper(): | ||
return | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r22992 | def check(func, connect=connectpipe): | ||
Yuya Nishihara
|
r40351 | stdout.flush() | ||
Yuya Nishihara
|
r22991 | server = connect() | ||
Yuya Nishihara
|
r22566 | try: | ||
return func(server) | ||||
finally: | ||||
server.stdin.close() | ||||
server.wait() | ||||
Yuya Nishihara
|
r40625 | |||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r40625 | def checkwith(connect=connectpipe, **kwargs): | ||
def wrap(func): | ||||
return check(func, lambda: connect(**kwargs)) | ||||
Augie Fackler
|
r43346 | |||
Yuya Nishihara
|
r40625 | return wrap | ||