refactor(cmd): context

iyear · Sep 4, 2022 · 51290f2 · 51290f2
1 parent e027a24
commit 51290f2
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 15 deletions.
diff --git a/cmd/dl/dl.go b/cmd/dl/dl.go
@@ -1,13 +1,10 @@
 package dl
 
 import (
- "context"
  "fmt"
  "github.com/iyear/tdl/app/dl/dlurl"
  "github.com/iyear/tdl/pkg/consts"
  "github.com/spf13/cobra"
- "os"
- "os/signal"
 )
 
 var (
@@ -27,9 +24,6 @@ var Cmd = &cobra.Command{
  Short: "Download what you want",
  Example: "tdl dl",
  RunE: func(cmd *cobra.Command, args []string) error {
- ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
- defer cancel()
-
  proxy, err := cmd.Flags().GetString("proxy")
  if err != nil {
  return err
@@ -42,7 +36,7 @@ var Cmd = &cobra.Command{
 
  switch mode {
  case consts.DownloadModeURL:
- if err := dlurl.Run(ctx, ns, proxy, partSize, threads, limit, urls); err != nil {
+ if err := dlurl.Run(cmd.Context(), ns, proxy, partSize, threads, limit, urls); err != nil {
  return fmt.Errorf("download failed: %v", err)
  }
  return nil

diff --git a/cmd/login/login.go b/cmd/login/login.go
@@ -1,22 +1,16 @@
 package login
 
 import (
- "context"
  "fmt"
  "github.com/iyear/tdl/app/login"
  "github.com/spf13/cobra"
- "os"
- "os/signal"
 )
 
 var Cmd = &cobra.Command{
  Use: "login",
  Short: "Login to Telegram",
  Example: "tdl login",
  RunE: func(cmd *cobra.Command, args []string) error {
- ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
- defer cancel()
-
  proxy, err := cmd.Flags().GetString("proxy")
  if err != nil {
  return err
@@ -26,7 +20,8 @@ var Cmd = &cobra.Command{
  if err != nil {
  return err
  }
- if err := login.Run(ctx, ns, proxy); err != nil {
+
+ if err := login.Run(cmd.Context(), ns, proxy); err != nil {
  return fmt.Errorf("login failed: %v", err)
  }
  return nil

diff --git a/cmd/root.go b/cmd/root.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+ "context"
  "fmt"
  "github.com/fatih/color"
  "github.com/iyear/tdl/cmd/dl"
@@ -10,6 +11,8 @@ import (
  "github.com/iyear/tdl/pkg/utils"
  "github.com/spf13/cobra"
  "github.com/spf13/cobra/doc"
+ "os"
+ "os/signal"
  "path/filepath"
 )
 
@@ -36,7 +39,10 @@ func init() {
 }
 
 func Execute() {
- if err := cmd.Execute(); err != nil {
+ ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)
+ defer cancel()
+
+ if err := cmd.ExecuteContext(ctx); err != nil {
  color.Red("%v", err)
  }
 }