feat: [python] identify basic import

phodal · Jan 15, 2020 · d08da6c · d08da6c
1 parent 6bd5da2
commit d08da6c
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 2 deletions.
diff --git a/pkg/domain/trial/code_file.go b/pkg/domain/trial/code_file.go
@@ -9,5 +9,11 @@ type CodeFile struct {
 	Members        []*CodeMember
 	DataStructures []CodeDataStruct
 	// Deprecated: support for migration only
-	ClassNodes     []domain.JClassNode
+	ClassNodes []domain.JClassNode
+}
+
+type CodeImport struct {
+	Source    string
+	AsName    string
+	UsageName []string
 }
diff --git a/trial/pkg/application/pyapp/py_ident_app_test.go b/trial/pkg/application/pyapp/py_ident_app_test.go
@@ -10,7 +10,6 @@ import (
 	"testing"
 )
 
-
 func Test_AllPythonGrammar(t *testing.T) {
 	g := NewGomegaWithT(t)
 
@@ -58,3 +57,15 @@ func Test_PythonClassWithDecorator(t *testing.T) {
 	g.Expect(codeFile.Members[0].MethodNodes[0].Name).To(Equal("bar"))
 	g.Expect(len(codeFile.Members[0].MethodNodes[0].Annotations.([]trial.PythonAnnotation))).To(Equal(2))
 }
+
+func Test_PythonImport(t *testing.T) {
+	g := NewGomegaWithT(t)
+
+	app := new(PythonApiApp)
+
+	file, _ := ioutil.ReadFile("testdata/grammar/import_stmt.py")
+	codeFile := app.Analysis(string(file), "import_stmt")
+
+	fmt.Println(codeFile)
+	g.Expect(true).To(Equal(true))
+}
diff --git a/trial/pkg/ast/pyast/python_ident_listener.go b/trial/pkg/ast/pyast/python_ident_listener.go
@@ -2,6 +2,7 @@ package pyast
 
 import (
 	"bytes"
+	"fmt"
 	"github.com/antlr/antlr4/runtime/Go/antlr"
 	parser "github.com/phodal/coca/languages/python"
 	"github.com/phodal/coca/pkg/domain/trial"
@@ -16,6 +17,7 @@ type PythonIdentListener struct {
 var currentCodeFile *trial.CodeFile
 var debug = false
 var output io.Writer
+var hasEnterMember = false
 
 func NewPythonIdentListener(fileName string) *PythonIdentListener {
 	currentCodeFile = &trial.CodeFile{}
@@ -32,7 +34,26 @@ func (s *PythonIdentListener) SetDebugOutput(isDebug bool) io.Writer {
 	return output
 }
 
+func (s *PythonIdentListener) EnterImport_stmt(ctx *parser.Import_stmtContext) {
+	for _, asName := range ctx.Dotted_as_names().(*parser.Dotted_as_namesContext).AllDotted_as_name() {
+		nameContext := asName.(*parser.Dotted_as_nameContext)
+		asNameText := nameContext.Dotted_name().GetText()
+		name := ""
+		if nameContext.Name() != nil {
+			name = nameContext.Name().GetText()
+		}
+
+		codeImport := &trial.CodeImport{
+			Source: asNameText,
+			AsName: name,
+		}
+
+		fmt.Println(codeImport)
+	}
+}
+
 func (s *PythonIdentListener) EnterClassdef(ctx *parser.ClassdefContext) {
+	hasEnterMember = true
 	dataStruct := trial.CodeDataStruct{
 		Name:       ctx.Name().GetText(),
 		ID:         "",
@@ -49,7 +70,12 @@ func (s *PythonIdentListener) EnterClassdef(ctx *parser.ClassdefContext) {
 	currentCodeFile.DataStructures = append(currentCodeFile.DataStructures, dataStruct)
 }
 
+func (s *PythonIdentListener) ExitClassdef(ctx *parser.ClassdefContext) {
+	hasEnterMember = false
+}
+
 func (s *PythonIdentListener) EnterFuncdef(ctx *parser.FuncdefContext) {
+	hasEnterMember = true
 	function := trial.CodeFunction{
 		Name: ctx.Name().GetText(),
 	}
@@ -68,6 +94,10 @@ func (s *PythonIdentListener) EnterFuncdef(ctx *parser.FuncdefContext) {
 	currentCodeFile.Members = append(currentCodeFile.Members, member)
 }
 
+func (s *PythonIdentListener) ExitFuncdef(ctx *parser.FuncdefContext) {
+	hasEnterMember = false
+}
+
 func BuildDecoratorsByIndex(node antlr.ParseTree, index int) []trial.PythonAnnotation {
 	var nodes []parser.DecoratorContext
 	for i := 0; i < index; i++ {