Merge pull request #576 from jwkvam/numpy-style

Python: add numpy docstring style.

UltiSnips/python.snippets

@@ -35,6 +35,7 @@ NORMAL  = 0x1
 DOXYGEN = 0x2
 SPHINX  = 0x3
 GOOGLE  = 0x4
+NUMPY   = 0x5
 
 SINGLE_QUOTES = "'"
 DOUBLE_QUOTES = '"'
@@ -105,6 +106,7 @@ def get_style(snip):
 	if    style == "doxygen": return DOXYGEN
 	elif  style == "sphinx": return SPHINX
 	elif  style == "google": return GOOGLE
+	elif  style == "numpy": return NUMPY
 	else: return NORMAL
 
 
@@ -117,6 +119,8 @@ def format_arg(arg, style):
 		return ":%s: TODO" % arg
 	elif style == GOOGLE:
 		return "%s (TODO): TODO" % arg
+	elif style == NUMPY:
+		return "%s : TODO" % arg
 
 
 def format_return(style):
@@ -139,6 +143,8 @@ def write_docstring_args(args, snip):
 
 	if style == GOOGLE:
 		write_google_docstring_args(args, snip)
+	elif style == NUMPY:
+		write_numpy_docstring_args(args, snip)
 	else:
 		for arg in args:
 			snip += format_arg(arg, style)
@@ -165,6 +171,23 @@ def write_google_docstring_args(args, snip):
 		snip.rv += '\n' + snip.mkline('', indent='')
 
 
+def write_numpy_docstring_args(args, snip):
+	if args:
+		snip += "Parameters"
+		snip += "----------"
+
+	kwargs = [arg for arg in args if arg.is_kwarg()]
+	args = [arg for arg in args if not arg.is_kwarg()]
+
+	if args:
+		for arg in args:
+			snip += format_arg(arg, NUMPY)
+	if kwargs:
+		for kwarg in kwargs:
+			snip += format_arg(kwarg, NUMPY) + ', optional'
+	snip.rv += '\n' + snip.mkline('', indent='')
+
+
 def write_init_body(args, parents, snip):
 	parents = [p.strip() for p in parents.split(",")]
 	parents = [p for p in parents if p != 'object']
@@ -199,7 +222,13 @@ def write_function_docstring(t, snip):
 		write_docstring_args(args, snip)
 
 	style = get_style(snip)
-	snip += format_return(style)
+
+	if style == NUMPY:
+		snip += 'Returns'
+		snip += '-------'
+		snip += 'TODO'
+	else:
+		snip += format_return(style)
 	snip.rv += '\n' + snip.mkline('', indent='')
 	snip += triple_quotes(snip)