Add doc string for models in sqlflow_models (#1717)

* Add doc string for models in sqlflow_models

* Fix unit tests

* Wait until tensorflow imported to fix unit test

* Wait longer

Author: shendiaomo

cmd/repl/repl_test.go

@@ -49,7 +49,7 @@ func testMainFastFail(t *testing.T, interactive bool) {
 
 	done := make(chan error)
 	go func() { done <- cmd.Wait() }()
-	timeout := time.After(2 * time.Second) // 2s are enough for **fast** fail
+	timeout := time.After(4 * time.Second) // 4s are enough for **fast** fail
 
 	select {
 	case <-timeout:
@@ -137,7 +137,7 @@ func TestComplete(t *testing.T) {
 
 	p.InsertText(`RAIN `, false, true)
 	c = s.completer(*p.Document())
-	a.Equal(11, len(c))
+	a.Equal(18, len(c))
 
 	p.InsertText(`DNN`, false, true)
 	c = s.completer(*p.Document())

pkg/sql/codegen/attribute/attribute.go

@@ -16,6 +16,8 @@ package attribute
 import (
 	"encoding/json"
 	"fmt"
+	"log"
+	"os/exec"
 	"reflect"
 	"sort"
 	"strings"
@@ -158,6 +160,18 @@ func NewDictionaryFromModelDefinition(estimator, prefix string) Dictionary {
 // PremadeModelParamsDocs stores parameters and documents of all known models
 var PremadeModelParamsDocs map[string]map[string]string
 
+// ExtractDocString extracts parameter documents from python doc strings
+func ExtractDocString(module ...string) {
+	cmd := exec.Command("python", "-uc", fmt.Sprintf("__import__('extract_docstring').print_param_doc('%s')", strings.Join(module, "', '")))
+	output, e := cmd.CombinedOutput()
+	if e != nil {
+		log.Println("ExtractDocString failed: ", e, string(output))
+	}
+	if e := json.Unmarshal(output, &PremadeModelParamsDocs); e != nil {
+		log.Println("ExtractDocString failed:", e, string(output))
+	}
+}
+
 func removeUnnecessaryParams() {
 	// The following parameters of canned estimators are already supported in the COLUMN clause.
 	for _, v := range PremadeModelParamsDocs {
@@ -171,5 +185,6 @@ func init() {
 	if err := json.Unmarshal([]byte(ModelParameterJSON), &PremadeModelParamsDocs); err != nil {
 		panic(err) // assertion
 	}
+	ExtractDocString("sqlflow_models")
 	removeUnnecessaryParams()
 }

pkg/sql/codegen/attribute/attribute_test.go

@@ -43,7 +43,7 @@ func TestDictionaryValidate(t *testing.T) {
 func TestPremadeModelParamsDocs(t *testing.T) {
 	a := assert.New(t)
 
-	a.Equal(11, len(PremadeModelParamsDocs))
+	a.Equal(18, len(PremadeModelParamsDocs))
 	a.Equal(len(PremadeModelParamsDocs["DNNClassifier"]), 12)
 	a.NotContains(PremadeModelParamsDocs["DNNClassifier"], "feature_columns")
 	a.Contains(PremadeModelParamsDocs["DNNClassifier"], "optimizer")

python/extract_docstring.py

@@ -74,6 +74,16 @@ def parse_ctor_args(f, prefix=''):
             [' '.join(doc.split()).replace("`", "'") for doc in total[2::2]]))
 
 
+def print_param_doc(*modules):
+    param_doc = {}  # { "class_names": {"parameters": "splitted docstrings"} }
+    for module in modules:
+        models = filter(lambda m: inspect.isclass(m[1]),
+                        inspect.getmembers(__import__(module)))
+        for name, cls in models:
+            param_doc[f'{module}.{name}'] = parse_ctor_args(cls, ':param')
+    print(json.dumps(param_doc))
+
+
 if __name__ == "__main__":
     param_doc = {}  # { "class_names": {"parameters": "splitted docstrings"} }
 

scripts/test/ipython.sh

@@ -45,6 +45,7 @@ DATASOURCE="mysql://root:root@tcp(127.0.0.1:3306)/?maxAllowedPacket=0"
 export PYTHONPATH=$GOPATH/src/sqlflow.org/sqlflow/python
 
 sqlflowserver &
+sleep 10
 # e2e test for standard SQL
 SQLFLOW_DATASOURCE=${DATASOURCE} SQLFLOW_SERVER=localhost:50051 ipython python/test_magic.py
 # TODO(yi): Re-enable the end-to-end test of Ant XGBoost after accelerating Travis CI.