adds must versions of module methods (#9)

dipdup-io · Sep 14, 2023 · 407520f · 407520f
1 parent 8c14ad3
commit 407520f
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 0 deletions.
diff --git a/examples/cron/custom.go b/examples/cron/custom.go
@@ -46,11 +46,29 @@ func (m *CustomModule) Input(name string) (*modules.Input, error) {
 	}
 }
 
+// MustInput -
+func (m *CustomModule) MustInput(name string) *modules.Input {
+	input, err := m.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 // Output -
 func (m *CustomModule) Output(name string) (*modules.Output, error) {
 	return nil, errors.Wrap(modules.ErrUnknownOutput, name)
 }
 
+// MustOutput -
+func (m *CustomModule) MustOutput(name string) *modules.Output {
+	output, err := m.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 // AttachTo -
 func (m *CustomModule) AttachTo(outputM modules.Module, outputName, inputName string) error {
 	output, err := outputM.Output(outputName)

diff --git a/examples/grpc/client.go b/examples/grpc/client.go
@@ -106,6 +106,15 @@ func (client *Client) Input(name string) (*modules.Input, error) {
 	return nil, errors.Wrap(modules.ErrUnknownInput, name)
 }
 
+// MustInput -
+func (client *Client) MustInput(name string) *modules.Input {
+	input, err := client.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 // Output -
 func (client *Client) Output(name string) (*modules.Output, error) {
 	if name != "time" {
@@ -114,6 +123,15 @@ func (client *Client) Output(name string) (*modules.Output, error) {
 	return client.output, nil
 }
 
+// MustOutput -
+func (client *Client) MustOutput(name string) *modules.Output {
+	output, err := client.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 // AttachTo -
 func (client *Client) AttachTo(outputModule modules.Module, outputName, inputName string) error {
 	outputChannel, err := outputModule.Output(outputName)

diff --git a/examples/grpc/custom.go b/examples/grpc/custom.go
@@ -39,11 +39,29 @@ func (m *CustomModule) Input(name string) (*modules.Input, error) {
 	return m.input, nil
 }
 
+// MustInput -
+func (m *CustomModule) MustInput(name string) *modules.Input {
+	input, err := m.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 // Output -
 func (m *CustomModule) Output(name string) (*modules.Output, error) {
 	return nil, errors.Wrap(modules.ErrUnknownOutput, name)
 }
 
+// MustOutput -
+func (m *CustomModule) MustOutput(name string) *modules.Output {
+	output, err := m.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 // AttachTo -
 func (m *CustomModule) AttachTo(outputModule modules.Module, outputName, inputName string) error {
 	outputChannel, err := outputModule.Output(outputName)

diff --git a/examples/zipper/custom.go b/examples/zipper/custom.go
@@ -34,6 +34,15 @@ func (m *CustomModule) Input(name string) (*modules.Input, error) {
 	return nil, errors.Wrap(modules.ErrUnknownInput, name)
 }
 
+// MustInput -
+func (m *CustomModule) MustInput(name string) *modules.Input {
+	input, err := m.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 // Output -
 func (m *CustomModule) Output(name string) (*modules.Output, error) {
 	if name != "output" {
@@ -42,6 +51,15 @@ func (m *CustomModule) Output(name string) (*modules.Output, error) {
 	return m.output, nil
 }
 
+// MustOutput -
+func (m *CustomModule) MustOutput(name string) *modules.Output {
+	output, err := m.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 // AttachTo -
 func (m *CustomModule) AttachTo(outputModule modules.Module, outputName, inputName string) error {
 	outputChannel, err := outputModule.Output(outputName)

diff --git a/pkg/modules/baseModule.go b/pkg/modules/baseModule.go
@@ -49,6 +49,14 @@ func (m *BaseModule) Input(name string) (*Input, error) {
 	return input, nil
 }
 
+func (m *BaseModule) MustInput(name string) *Input {
+	input, err := m.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 func (m *BaseModule) CreateInput(name string) {
 	m.inputs.Set(name, NewInput(name))
 }
@@ -61,6 +69,14 @@ func (m *BaseModule) Output(name string) (*Output, error) {
 	return output, nil
 }
 
+func (m *BaseModule) MustOutput(name string) *Output {
+	output, err := m.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 func (m *BaseModule) CreateOutput(name string) {
 	m.outputs.Set(name, NewOutput(name))
 }

diff --git a/pkg/modules/module.go b/pkg/modules/module.go
@@ -21,6 +21,8 @@ type Module interface {
 	Start(ctx context.Context)
 
 	Input(name string) (*Input, error)
+	MustInput(name string) *Input
 	Output(name string) (*Output, error)
+	MustOutput(name string) *Output
 	AttachTo(output Module, outputName, inputName string) error
 }
diff --git a/pkg/modules/zipper/module.go b/pkg/modules/zipper/module.go
@@ -70,6 +70,15 @@ func (m *Module[Key]) Input(name string) (*modules.Input, error) {
 	}
 }
 
+// MustInput - returns input by name
+func (m *Module[Key]) MustInput(name string) *modules.Input {
+	input, err := m.Input(name)
+	if err != nil {
+		panic(err)
+	}
+	return input
+}
+
 // Output - returns output by name
 func (m *Module[Key]) Output(name string) (*modules.Output, error) {
 	if name != OutputName {
@@ -78,6 +87,15 @@ func (m *Module[Key]) Output(name string) (*modules.Output, error) {
 	return m.output, nil
 }
 
+// MustOutput - returns output by name
+func (m *Module[Key]) MustOutput(name string) *modules.Output {
+	output, err := m.Output(name)
+	if err != nil {
+		panic(err)
+	}
+	return output
+}
+
 // AttachTo - attach input to output with name
 func (m *Module[Key]) AttachTo(outputModule modules.Module, outputName, inputName string) error {
 	outputChannel, err := outputModule.Output(outputName)